feature: beacon on commandline

README.md

@@ -106,13 +106,17 @@ and its [VirusTotal result](https://github.com/dobin/ShellcodeObfuscationLab/blo
 
 Use the `x64 native tools command prompt` from Visual Studio so you have access to `ml.exe` and `Windows.h`. 
 
-To compile all the source from `chromatophore/` into `output/*.exe`:
+To compile all the source from `chromatophore/` into `output/*.exe` (uses default shellcode file `./beacon.bin`):
 ```
 > python.exe sol.py compile
 ```
 
-To send all the exes to VirusTotal:
+To compile with a custom shellcode file:
+```
+> python.exe sol.py compile --shellcode custom_payload.bin
+```
 
+To send all the exes to VirusTotal:
 ```
 > set VT_API_KEY=123...
 > python.exe sol.py vt

sol.py

@@ -1,6 +1,7 @@
 import subprocess
 import os
 import sys
+import argparse
 
 from vt import scan_files
 
@@ -56,14 +57,13 @@ function_map = {
 }
 
 
-def compile_all():
+def compile_all(shellcode_file="beacon.bin"):
     clean_files()
 
     for module in function_map.keys():
         print("Templating")
         print("Module: " + module)
 
-        shellcode_file = "beacon.bin"
         mod_data = function_map[module](shellcode_file)
 
         template_input = "chromatophore\\{}\\{}.c".format(module, module)
@@ -73,10 +73,9 @@ def compile_all():
         compile_module(module)
 
 
-def test_module():
+def test_module(shellcode_file="beacon.bin"):
     module = "noobfuscation"
 
-    shellcode_file = "beacon.bin"
     mod_data = function_map[module](shellcode_file)
     print("Mod data: " + mod_data)
 
@@ -147,16 +146,24 @@ def execute_module(module):
 
 
 def main():
-    if sys.argv[1] == "clean":
+    parser = argparse.ArgumentParser(description='Malware development encryption tool')
+    parser.add_argument('command', choices=['clean', 'compile', 'vt', 'test'], 
+                       help='Command to execute')
+    parser.add_argument('-s', '--shellcode', default='beacon.bin',
+                       help='Shellcode file to use (default: beacon.bin)')
+    
+    args = parser.parse_args()
+    
+    if args.command == "clean":
         clean_files()
-    elif sys.argv[1] == "compile":
-        compile_all()
-    elif sys.argv[1] == "vt":
+    elif args.command == "compile":
+        compile_all(args.shellcode)
+    elif args.command == "vt":
         scan_files()
-    elif sys.argv[1] == "test":
-        test_module()
+    elif args.command == "test":
+        test_module(args.shellcode)
     else:
-        print("Invalid argument. Use 'clean', 'compile', or 'vt'.")
+        print("Invalid argument. Use 'clean', 'compile', 'vt', or 'test'.")
 
 if __name__ == "__main__":
     main()