Update __main__.py

balsamic/__main__.py

@@ -1,28 +1,25 @@
 from balsamic import balsamic
-import argparse
 if __name__ == "__main__":
-	#import argparse
+    # Importing argparse
     import argparse
 
-	#defines parser function
+    # Define the parser function
     def parser():
         
-		#create base argument parser
+        # Create the base argument parser
         parser = argparse.ArgumentParser(description="balsamic args")
         subparse = parser.add_subparsers(dest="attack")
 
-		#create subparser for webreq attack
+        # Create subparsers for different attacks
         webreqparser = subparse.add_parser("webreq")
-		webreqparser.add_argument("-s","--schema",required=True, choices=["http","https"])
         webreqparser.add_argument("-m", "--method")
-		webreqparser.add_argument("-rh","--rhost",required=True)
-		webreqparser.add_argument("-rp","--rport",required=True)
+        webreqparser.add_argument("-u", "--url", required=True)
         webreqparser.add_argument("-p", "--parameter")
         webreqparser.add_argument("-co", "--cookie")
         webreqparser.add_argument("-P", "--payload", required=True)
         webreqparser.add_argument("-c", "--command")
+        webreqparser.add_argument("-H", "--headers")  # New argument for custom headers
 
-		#create subparser for socksend attack
         socksendparser = subparse.add_parser("socksend")
         socksendparser.add_argument("-rh", "--rhost", required=True)
         socksendparser.add_argument("-rp", "--rport", required=True)
@@ -31,7 +28,6 @@ if __name__ == "__main__":
         socksendparser.add_argument("-s", "--steps", default="0")
         socksendparser.add_argument("-e", "--encode", action="store_true")
 
-		#create subparser for socklisten attack
         socklistenparser = subparse.add_parser("socklisten")
         socklistenparser.add_argument("-lp", "--lport", required=True)
         socklistenparser.add_argument("-P", "--payload", required=True)
@@ -39,7 +35,7 @@ if __name__ == "__main__":
         socklistenparser.add_argument("-s", "--steps", default="0")
         socklistenparser.add_argument("-e", "--encode", action="store_true")
         
-		#return parsed arguments
+        # Return parsed arguments
         args = parser.parse_args()
         return args
 
@@ -47,8 +43,13 @@ if __name__ == "__main__":
     if args.command:
         balsamic.updatecmd(args.command)
 
+    # Extract custom headers from the arguments
+    try:
+        custom_headers = dict([header.split(":") for header in args.headers.split(";")]) if args.headers else None
+    except AttributeError:
+        pass
     if args.attack == "webreq":
-		balsamic.webreq(args.schema,args.method,args.rhost,args.rport,args.payload,args.parameter,args.cookie)
+        balsamic.webreq(args.method, args.url, args.payload, args.parameter, args.cookie, custom_headers=custom_headers)
     elif args.attack == "socksend":
         balsamic.socksend(args.rhost, args.rport, args.payload, args.encode, args.steps)
     elif args.attack == "socklisten":