mirror of
https://gitlab.sectorq.eu/jaydee/api_server.git
synced 2026-09-17 02:12:55 +02:00
139 lines
4.7 KiB
Python
139 lines
4.7 KiB
Python
|
|
import requests
|
|
import sys
|
|
from prompt_toolkit.shortcuts import checkboxlist_dialog
|
|
from prompt_toolkit.shortcuts import radiolist_dialog
|
|
from prompt_toolkit import prompt
|
|
import logging
|
|
import json
|
|
import urllib3
|
|
import re
|
|
import os
|
|
import hvac
|
|
|
|
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s [%(levelname)s] %(message)s"
|
|
)
|
|
|
|
class LoginError(Exception):
|
|
"""Raised when login fails due to invalid credentials or missing CSRF token."""
|
|
|
|
|
|
def setup_vault():
|
|
"""Sets up and returns a Vault client."""
|
|
vault_addr = os.environ.get("VAULT_ADDR", "https://vault.sectorq.eu")
|
|
try:
|
|
value_token = os.environ.get("VAULT_TOKEN")
|
|
if value_token is None:
|
|
raise KeyError
|
|
except KeyError:
|
|
value_token = prompt("Valult root token : ", is_password=True)
|
|
print("Use command export VAULT_TOKEN=")
|
|
os.environ["VAULT_TOKEN"] = value_token
|
|
client = hvac.Client(url=vault_addr, token=value_token)
|
|
|
|
# Check if connected
|
|
if client.is_authenticated():
|
|
print("Connected to Vault")
|
|
else:
|
|
print("Failed to login")
|
|
sys.exit(1)
|
|
return client
|
|
|
|
|
|
def get_secret(client, path, field=None):
|
|
"""Retrieves a secret from Vault at the specified path."""
|
|
return client.secrets.kv.v2.read_secret_version(
|
|
path=path, mount_point="secret", raise_on_deleted_version=True
|
|
)["data"]["data"][field]
|
|
|
|
def get_token(args):
|
|
"""Get CSRF token and cookies"""
|
|
vault_client = setup_vault()
|
|
unifi_password = get_secret(vault_client, "unifi/password", field="value")
|
|
print("Logging in to Unifi Controller...")
|
|
session = requests.Session()
|
|
login_url = "https://unifi.sectorq.eu/api/auth/login"
|
|
payload_login = {"username": "jaydee", "password": unifi_password}
|
|
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)
|
|
logging.info(f"received payload: {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 = re.split(r'[, :]+', args.ips)
|
|
ips = list(dict.fromkeys(ips + ips_to_add))
|
|
ips = [x for x in ips if x]
|
|
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
|
|
)
|
|
logging.info(f"IP Banned {ips}" if args.action == "ban_ip" else f"IP Unbanned {ips}")
|
|
return ("Update response status:", update_response.text)
|
|
|
|
|
|
def run(data):
|
|
try:
|
|
return ban_ip(data)
|
|
except Exception as e:
|
|
logging.error(f"Error occurred while processing {data.action} for IP {data.ips}: {e}")
|
|
return f"{data.action} requested {data.ips}"
|
|
|