Merge pull request #13 from mmstick/patch-1

Add compression parameters

1 files changed, 69 insertions, 38 deletions

diff --git a/main.go b/main.go
index 0461b68..40fb246 100644
--- a/main.go
+++ b/main.go
@@ -1,67 +1,98 @@
 package main
 
 import (
-	"flag"
+	"github.com/pwaller/goupx/hemfix"
 	"log"
 	"os"
 	"os/exec"
-
-	"github.com/pwaller/goupx/hemfix"
 )
 
-const usageText = "usage: goupx [args...] path\n"
+const usageText = `usage: goupx [args...] files...
+      
+    --no-upx: Disables UPX from running.
+    --strip-binary: Strips binaries before compressing them.
+      
+See UPX's documentation (man upx) for information on UPX's flags.
+`
+
+var run_strip = false
+var run_upx = true
 
-// usage prints some nice output instead of panic stacktrace when an user calls goupx without arguments
+// 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() {
-
-	run_strip := flag.Bool("s", false, "run strip")
-	run_upx := flag.Bool("u", true, "run upx")
-
-	flag.Parse()
-
-	if flag.NArg() != 1 {
+// parseArguments parses arguments from os.Args and separates the goupx flags
+// from the UPX flags, as well as separating the files from the arguments.
+func parseArguments() (args []string, files []string) {
+	if len(os.Args) == 1 {
 		usage()
-		return
 	}
-
-	defer func() {
-		if err := recover(); err != nil {
-			log.Print("Panicked. Giving up.")
-			panic(err)
-			return
+	args = append(args, "/usr/bin/upx")
+	for _, arg := range os.Args[1:] {
+		switch {
+		case arg == "-h" || arg == "--help":
+			usage()
+		case arg == "--no-upx":
+			run_upx = false
+		case arg == "--strip-binary":
+			run_strip = true
+		case arg[0] != '-':
+			files = append(files, arg)
+		default:
+			args = append(args, arg)
 		}
-	}()
-
-	input_file := flag.Arg(0)
-	err := hemfix.FixFile(input_file)
-	if err != nil {
-		log.Panicf("Failed to fix '%s': %v", input_file, err)
 	}
-	log.Print("File fixed!")
+	return
+}
 
-	if *run_strip {
-		cmd := exec.Command("strip", "-s", input_file)
+// compressBinary attempts to compress the binary with UPX.
+func compressBinary(input_file string, arguments []string) {
+	if run_upx {
+		cmd := &exec.Cmd{
+			Path: "/usr/bin/upx",
+			Args: append(arguments, input_file),
+		}
 		cmd.Stdout = os.Stdout
 		cmd.Stderr = os.Stderr
-		err = cmd.Run()
-		if err != nil {
-			log.Panic("strip failed: ", err)
+		if err := cmd.Run(); err != nil {
+			log.Panic("upx failed: ", err)
 		}
 	}
+}
 
-	if *run_upx {
-		cmd := exec.Command("upx", input_file)
+// stripBinary attempts to strip the binary.
+func stripBinary(input_file string) {
+	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("upx failed: ", err)
+		if err := cmd.Run(); err != nil {
+			log.Panic("strip failed: ", err)
 		}
 	}
+}
+
+// runHemfix will attempt to fix the current input file.
+func runHemfix(input_file string) {
+	if err := hemfix.FixFile(input_file); err != nil {
+		log.Panicf("Failed to fix '%s': %v", input_file, err)
+	}
+	log.Print("File fixed!")
+}
 
+func main() {
+	arguments, files := parseArguments()
+	for _, file := range files {
+		runHemfix(file)
+		stripBinary(file)
+		compressBinary(file, arguments)
+	}
+	if err := recover(); err != nil {
+		log.Print("Panicked. Giving up.")
+		panic(err)
+		return
+	}
 }