diff --git a/app/main.py b/app/main.py index f422a23..66a633b 100644 --- a/app/main.py +++ b/app/main.py @@ -6,8 +6,8 @@ from app.scripts import script1, script2 app = FastAPI() class Payload(BaseModel): - name: str action: str + ips: str @app.get("/") def root(): diff --git a/app/scripts/script1.py b/app/scripts/script1.py index 8533faa..d2e9a9f 100644 --- a/app/scripts/script1.py +++ b/app/scripts/script1.py @@ -1,3 +1,91 @@ -def run(data): - return f"{data.name} requested {data.action}" \ No newline at end of file +import requests +import sys +from prompt_toolkit.shortcuts import checkboxlist_dialog +from prompt_toolkit.shortcuts import radiolist_dialog +from prompt_toolkit import prompt +import json +class LoginError(Exception): + """Raised when login fails due to invalid credentials or missing CSRF token.""" + +def get_token(args): + """Get CSRF token and cookies""" + session = requests.Session() + login_url = "https://unifi.sectorq.eu/api/auth/login" + payload_login = {"username": "jaydee", "password": "l4c1j4yd33Du5lo"} + response = session.post(login_url, json=payload_login, verify=False) + response.raise_for_status() + + # Extract CSRF token from headers + csrf_token = response.headers.get("x-csrf-token") + cookies = session.cookies.get_dict() + if not csrf_token: + raise LoginError("CSRF token not found, login failed") + return csrf_token, cookies, session + +def ban_ip(args): + """Ban or unban IP address""" + data = get_token(args) + + headers = {"x-csrf-token": data[0]} + + # Step 3: Update firewall group + # print(csrf_token) + get_url = "https://unifi.sectorq.eu/proxy/network/api/s/default/rest/firewallgroup" + get_response = data[2].get(get_url, cookies=data[1], headers=headers, verify=False) + ips = next( + ( + item["group_members"] + for item in json.loads(get_response.text)["data"] + if item["name"] == "file2ban" + ), + None, + ) + if args.action == "unban_ip": + commands_tuples = [(cmd, cmd) for cmd in sorted(ips)] + commands_tuples.insert(0, ("__ALL__", "[Select ALL]")) + value_in = checkboxlist_dialog( + title="Select Services", + text="Choose one or more services:", + values=commands_tuples, + ).run() + ips = [x for x in ips if x not in value_in] + + elif args.action == "ban_ip": + if args.ips is None: + ip = input("Enter IP to ban: ") + if not ip: + sys.exit("No IP provided") + ips.append(ip) + else: + ips_to_add = [] + if args.ips == "NAS": + with open("/etc/config/ipsec_deny.conf", "r", encoding="utf-8") as f: + for line in f: + ip = line.strip().split(":")[1] + ips_to_add.append(ip) + + else: + ips_to_add = args.ips.split(",") + ips = list(dict.fromkeys(ips + ips_to_add)) + + update_url = f"https://unifi.sectorq.eu/proxy/network/api/s/default/rest/firewallgroup/6782dbc2ffa79454f61fac43" + payload_update = { + "name": "file2ban", + "group_type": "address-group", + "group_members": ips, + "site_id": "site-id", # Replace with your actual site ID + "_id": "6782dbc2ffa79454f61fac43", # Replace with your actual firewall group ID + } + update_response = data[2].put( + update_url, json=payload_update, headers=headers, verify=False + ) + print("Update response status:", update_response.text) + + +def run(data): + ban_ip(data) + + + return f"{data.action} requested {data.ip}" +