Pass Arguments Directly

Replaced the flag package with a custom implementation utilizing os.Args.
I've also fixed the program so that it can compress multiple Go binaries.

1 files changed, 55 insertions, 68 deletions

diff --git a/main.go b/main.go
index 8530940..9aec87c 100644
--- a/main.go
+++ b/main.go
@@ -1,8 +1,6 @@
 package main
 
 import (
-	"flag"
-	"fmt"
 	"github.com/pwaller/goupx/hemfix"
 	"log"
 	"os"
@@ -11,92 +9,81 @@ import (
 
 const usageText = "usage: goupx [args...] path\n"
 
-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")
-)
+var run_strip = false
+var run_upx = true
 
 // usage prints some nice output instead of panic stacktrace when an user calls
 // goupx without arguments
 func usage() {
 	os.Stderr.WriteString(usageText)
-	flag.PrintDefaults()
-}
-
-// 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)
-	}
 }
 
-// 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)
+// 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) {
+	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)
+		}
 	}
+	return
 }
 
 // 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 compressBinary(input_file string, arguments []string) {
+	if run_upx {
+		cmd := &exec.Cmd{
+			Path:   "/usr/bin/upx",
+			Args:   append(arguments, input_file),
+			Stdout: os.Stdout,
+			Stderr: os.Stderr,
+		}
+		if err := cmd.Run(); err != nil {
+			log.Panic("upx failed: ", err)
+		}
 	}
 }
 
-func main() {
-	flag.Parse()
-
-	if flag.NArg() != 1 {
-		usage()
-		return
-	}
-
-	checkCompressionLevel()
-
-	defer func() {
-		if err := recover(); err != nil {
-			log.Print("Panicked. Giving up.")
-			panic(err)
-			return
+// 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
+		if err := cmd.Run(); err != nil {
+			log.Panic("strip failed: ", err)
 		}
-	}()
+	}
+}
 
-	input_file := flag.Arg(0)
-	err := hemfix.FixFile(input_file)
-	if err != nil {
+// 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!")
+}
 
-	if *run_strip {
-		stripBinary(input_file)
+func main() {
+	arguments, files := parseArguments()
+	for _, file := range files {
+		runHemfix(file)
+		stripBinary(file)
+		compressBinary(file, arguments)
 	}
-
-	if *run_upx {
-		compressBinary(input_file)
+	if err := recover(); err != nil {
+		log.Print("Panicked. Giving up.")
+		panic(err)
+		return
 	}
 }