From 96ef91b5f95ba600587c3b3f633f6936651106d8 Mon Sep 17 00:00:00 2001
From: Michael Hansen <zrax0111@gmail.com>
Date: Tue, 26 Apr 2022 14:44:10 -0700
Subject: [PATCH] Allow specifying an output filename on the command line,
 instead of forcing reliance on shell redirection.

Closes #31
---
 pycdas.cpp              | 41 ++++++++++++++++++++++++++++++++++-------
 pycdc.cpp               | 41 ++++++++++++++++++++++++++++++++++-------
 tests/decompyle_test.sh |  2 +-
 3 files changed, 69 insertions(+), 15 deletions(-)

diff --git a/pycdas.cpp b/pycdas.cpp
index bb14665..1bf4288 100644
--- a/pycdas.cpp
+++ b/pycdas.cpp
@@ -59,7 +59,7 @@ static void ivprintf(int indent, const char* fmt, va_list varargs)
 {
     for (int i=0; i<indent; i++)
         fputs("    ", pyc_output);
-    vprintf(fmt, varargs);
+    vfprintf(pyc_output, fmt, varargs);
 }
 
 static void iprintf(int indent, const char* fmt, ...)
@@ -231,26 +231,53 @@ void output_object(PycRef<PycObject> obj, PycModule* mod, int indent)
 
 int main(int argc, char* argv[])
 {
-    if (argc < 2) {
+    const char* infile = nullptr;
+    for (int arg = 1; arg < argc; ++arg) {
+        if (strcmp(argv[arg], "-o") == 0) {
+            if (arg + 1 < argc) {
+                const char* filename = argv[++arg];
+                FILE* outfile = fopen(filename, "w");
+                if (!outfile) {
+                    fprintf(stderr, "Error opening file '%s' for writing\n",
+                            argv[arg]);
+                    return 1;
+                }
+                pyc_output = outfile;
+            } else {
+                fputs("Option '-o' requires a filename\n", stderr);
+                return 1;
+            }
+        } else if (strcmp(argv[arg], "--help") == 0 || strcmp(argv[arg], "-h") == 0) {
+            fprintf(stderr, "Usage:  %s [options] input.pyc\n\n", argv[0]);
+            fputs("Options:\n", stderr);
+            fputs("  -o <filename>   Write output to <filename> (default: stdout)\n", stderr);
+            fputs("  --help          Show this help text and then exit\n", stderr);
+            return 0;
+        } else {
+            infile = argv[arg];
+        }
+    }
+
+    if (!infile) {
         fputs("No input file specified\n", stderr);
         return 1;
     }
 
     PycModule mod;
     try {
-        mod.loadFromFile(argv[1]);
+        mod.loadFromFile(infile);
     } catch (std::exception& ex) {
-        fprintf(stderr, "Error disassembling %s: %s\n", argv[1], ex.what());
+        fprintf(stderr, "Error disassembling %s: %s\n", infile, ex.what());
         return 1;
     }
-    const char* dispname = strrchr(argv[1], PATHSEP);
-    dispname = (dispname == NULL) ? argv[1] : dispname + 1;
+    const char* dispname = strrchr(infile, PATHSEP);
+    dispname = (dispname == NULL) ? infile : dispname + 1;
     fprintf(pyc_output, "%s (Python %d.%d%s)\n", dispname, mod.majorVer(), mod.minorVer(),
            (mod.majorVer() < 3 && mod.isUnicode()) ? " -U" : "");
     try {
         output_object(mod.code().cast<PycObject>(), &mod, 0);
     } catch (std::exception& ex) {
-        fprintf(stderr, "Error disassembling %s: %s\n", argv[1], ex.what());
+        fprintf(stderr, "Error disassembling %s: %s\n", infile, ex.what());
         return 1;
     }
 
diff --git a/pycdc.cpp b/pycdc.cpp
index 4d8d7fa..97a0eae 100644
--- a/pycdc.cpp
+++ b/pycdc.cpp
@@ -9,31 +9,58 @@
 
 int main(int argc, char* argv[])
 {
-    if (argc < 2) {
+    const char* infile = nullptr;
+    for (int arg = 1; arg < argc; ++arg) {
+        if (strcmp(argv[arg], "-o") == 0) {
+            if (arg + 1 < argc) {
+                const char* filename = argv[++arg];
+                FILE* outfile = fopen(filename, "w");
+                if (!outfile) {
+                    fprintf(stderr, "Error opening file '%s' for writing\n",
+                            argv[arg]);
+                    return 1;
+                }
+                pyc_output = outfile;
+            } else {
+                fputs("Option '-o' requires a filename\n", stderr);
+                return 1;
+            }
+        } else if (strcmp(argv[arg], "--help") == 0 || strcmp(argv[arg], "-h") == 0) {
+            fprintf(stderr, "Usage:  %s [options] input.pyc\n\n", argv[0]);
+            fputs("Options:\n", stderr);
+            fputs("  -o <filename>   Write output to <filename> (default: stdout)\n", stderr);
+            fputs("  --help          Show this help text and then exit\n", stderr);
+            return 0;
+        } else {
+            infile = argv[arg];
+        }
+    }
+
+    if (!infile) {
         fputs("No input file specified\n", stderr);
         return 1;
     }
 
     PycModule mod;
     try {
-        mod.loadFromFile(argv[1]);
+        mod.loadFromFile(infile);
     } catch (std::exception& ex) {
-        fprintf(stderr, "Error loading file %s: %s\n", argv[1], ex.what());
+        fprintf(stderr, "Error loading file %s: %s\n", infile, ex.what());
         return 1;
     }
     if (!mod.isValid()) {
-        fprintf(stderr, "Could not load file %s\n", argv[1]);
+        fprintf(stderr, "Could not load file %s\n", infile);
         return 1;
     }
-    const char* dispname = strrchr(argv[1], PATHSEP);
-    dispname = (dispname == NULL) ? argv[1] : dispname + 1;
+    const char* dispname = strrchr(infile, PATHSEP);
+    dispname = (dispname == NULL) ? infile : dispname + 1;
     fputs("# Source Generated with Decompyle++\n", pyc_output);
     fprintf(pyc_output, "# File: %s (Python %d.%d%s)\n\n", dispname, mod.majorVer(), mod.minorVer(),
             (mod.majorVer() < 3 && mod.isUnicode()) ? " Unicode" : "");
     try {
         decompyle(mod.code(), &mod);
     } catch (std::exception& ex) {
-        fprintf(stderr, "Error decompyling %s: %s\n", argv[1], ex.what());
+        fprintf(stderr, "Error decompyling %s: %s\n", infile, ex.what());
         return 1;
     }
 
diff --git a/tests/decompyle_test.sh b/tests/decompyle_test.sh
index 1a8f54c..2be87f9 100755
--- a/tests/decompyle_test.sh
+++ b/tests/decompyle_test.sh
@@ -41,7 +41,7 @@ upfiles=()
 for pyc in "${compfiles[@]}" "${xfcfiles[@]}"; do
     base="$outdir/$(basename "$pyc")"
 
-    ./pycdc "$pyc" 2>"$base.err" 1>"$base.src.py"
+    ./pycdc "$pyc" -o "$base.src.py" >"$base.err" 2>&1
     if (( $? )) || [[ -s "$base.err" ]]
     then
         if [[ "$(dirname "$pyc")" =~ xfail ]]