Add compression parameters

Added manual compression, best, brute and ultra-brute compression flags.
Manual compression is called with the -c flag and ranges from 1 through 9.
-best, -brute and -ultra-brute call their repective UPX compression levels.

1 files changed, 56 insertions, 21 deletions

diff --git a/main.go b/main.go
index 0461b68..8530940 100644
--- a/main.go
+++ b/main.go
@@ -2,26 +2,72 @@ package main
 
 import (
 	"flag"
+	"fmt"
+	"github.com/pwaller/goupx/hemfix"
 	"log"
 	"os"
 	"os/exec"
-
-	"github.com/pwaller/goupx/hemfix"
 )
 
 const usageText = "usage: goupx [args...] path\n"
 
-// usage prints some nice output instead of panic stacktrace when an user calls goupx without arguments
+var (
+	run_strip               = flag.Bool("s", false, "run strip")
+	run_upx                 = flag.Bool("u", true, "run upx")
+	manual_compression      = flag.Uint("c", 0, "Manual Compression Level: (1-9)")
+	best_compression        = flag.Bool("best", false, "Best Compression")
+	brute_compression       = flag.Bool("brute", false, "Brute Compression")
+	ultra_brute_compression = flag.Bool("ultra-brute", false, "Ultra Brute Compression")
+)
+
+// usage prints some nice output instead of panic stacktrace when an user calls
+// goupx without arguments
 func usage() {
 	os.Stderr.WriteString(usageText)
 	flag.PrintDefaults()
 }
 
-func main() {
+// checkCompressionLevel sets the compression level to 9 if it is higher than
+// the maximum limit
+func checkCompressionLevel() {
+	if *manual_compression > uint(9) {
+		*manual_compression = uint(9)
+	}
+}
 
-	run_strip := flag.Bool("s", false, "run strip")
-	run_upx := flag.Bool("u", true, "run upx")
+// stripBinary attempts to strip the binary.
+func stripBinary(input_file string) {
+	cmd := exec.Command("strip", "-s", input_file)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	if err := cmd.Run(); err != nil {
+		log.Panic("strip failed: ", err)
+	}
+}
 
+// compressBinary attempts to compress the binary with UPX.
+func compressBinary(input_file string) {
+	var cmd *exec.Cmd
+	if *best_compression {
+		cmd = exec.Command("upx", "--best", input_file)
+	} else if *brute_compression {
+		cmd = exec.Command("upx", "--brute", input_file)
+	} else if *ultra_brute_compression {
+		cmd = exec.Command("upx", "--ultra-brute", input_file)
+	} else if *manual_compression != 0 {
+		cmd = exec.Command("upx",
+			fmt.Sprintf("%s%d", "-", *manual_compression), input_file)
+	} else {
+		cmd = exec.Command("upx", input_file)
+	}
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	if err := cmd.Run(); err != nil {
+		log.Panic("upx failed: ", err)
+	}
+}
+
+func main() {
 	flag.Parse()
 
 	if flag.NArg() != 1 {
@@ -29,6 +75,8 @@ func main() {
 		return
 	}
 
+	checkCompressionLevel()
+
 	defer func() {
 		if err := recover(); err != nil {
 			log.Print("Panicked. Giving up.")
@@ -45,23 +93,10 @@ func main() {
 	log.Print("File fixed!")
 
 	if *run_strip {
-		cmd := exec.Command("strip", "-s", input_file)
-		cmd.Stdout = os.Stdout
-		cmd.Stderr = os.Stderr
-		err = cmd.Run()
-		if err != nil {
-			log.Panic("strip failed: ", err)
-		}
+		stripBinary(input_file)
 	}
 
 	if *run_upx {
-		cmd := exec.Command("upx", input_file)
-		cmd.Stdout = os.Stdout
-		cmd.Stderr = os.Stderr
-		err = cmd.Run()
-		if err != nil {
-			log.Panic("upx failed: ", err)
-		}
+		compressBinary(input_file)
 	}
-
 }