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
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 get_token(args):
    """Get CSRF token and cookies"""
    print("Logging in to Unifi Controller...")
    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)
    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))

    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}"