#!/usr/bin/env python3
import base64
import hashlib
import json
import math
import os
import re
import secrets
import ssl
import subprocess
import sys
from urllib.parse import quote_plus

import requests
from Crypto.Cipher import AES, PKCS1_v1_5
from Crypto.PublicKey import RSA
from Crypto.Util.Padding import pad, unpad
from requests import RequestException


AES_KEY_BYTES = 16
MIN_AES_KEY = 10 ** (AES_KEY_BYTES - 1)
MAX_AES_KEY = (10**AES_KEY_BYTES) - 1
PKCS1_V1_5_HEADER_BYTES = 11


def load_env_file():
    path = os.path.join(os.path.dirname(os.path.dirname(__file__)), ".env")
    try:
        with open(path, "r", encoding="utf-8") as env_file:
            for line in env_file:
                line = line.strip()
                if not line or line.startswith("#") or "=" not in line:
                    continue
                key, value = line.split("=", 1)
                os.environ.setdefault(key, value.strip().strip("\"'"))
    except FileNotFoundError:
        pass


def default_gateway():
    try:
        result = subprocess.run(
            ["/sbin/route", "-n", "get", "default"],
            check=True,
            capture_output=True,
            text=True,
            timeout=5,
        )
    except Exception:
        return None

    for line in result.stdout.splitlines():
        key, _, value = line.partition(":")
        if key.strip() == "gateway":
            return value.strip()
    return None


def config():
    load_env_file()
    host = os.environ.get("DECO_HOST") or default_gateway()
    password = os.environ.get("DECO_PASSWORD")
    if not password:
        raise RuntimeError("DECO_PASSWORD must be set in .env.")
    if not host:
        raise RuntimeError("DECO_HOST must be set in .env; default gateway detection failed.")
    if not host.startswith(("http://", "https://")):
        host = f"http://{host}"
    verify_ssl = os.environ.get("DECO_VERIFY_SSL", "true").lower() not in {
        "0",
        "false",
        "no",
    }
    return {
        "host": host.rstrip("/"),
        "username": os.environ.get("DECO_USERNAME", "admin"),
        "password": password,
        "verify_ssl": verify_ssl,
        "timeout": int(os.environ.get("DECO_TIMEOUT", "10")),
    }


def byte_len(n):
    return (int(math.log2(n)) + 8) >> 3


def rsa_encrypt(n, e, plaintext):
    public_key = RSA.construct((n, e)).publickey()
    encryptor = PKCS1_v1_5.new(public_key)
    block_size = byte_len(n)
    bytes_per_block = block_size - PKCS1_V1_5_HEADER_BYTES
    encrypted_text = ""
    for index in range(0, len(plaintext), bytes_per_block):
        encrypted_text += encryptor.encrypt(plaintext[index:index + bytes_per_block]).hex()
    return encrypted_text


def decode_name(value):
    if not value:
        return value
    try:
        return base64.b64decode(value).decode()
    except Exception:
        return value


def title_from_snake(value):
    if not value:
        return value
    return " ".join(part.title() for part in value.split("_"))


def deco_name(device):
    return (
        decode_name(device.get("custom_nickname"))
        or title_from_snake(device.get("nickname"))
        or device.get("device_model")
        or device.get("mac")
    )


class DecoApi:
    def __init__(self, host, username, password, verify_ssl, timeout):
        self.host = host
        self.username = username
        self.password = password
        self.verify_ssl = verify_ssl
        self.timeout = timeout
        self.session = requests.Session()
        self.aes_key = None
        self.aes_iv = None
        self.password_rsa_n = None
        self.password_rsa_e = None
        self.sign_rsa_n = None
        self.sign_rsa_e = None
        self.seq = None
        self.stok = None
        self.cookie = None

    def generate_aes_key_and_iv(self):
        self.aes_key = str(secrets.randbelow(MAX_AES_KEY - MIN_AES_KEY) + MIN_AES_KEY).encode()
        self.aes_iv = str(secrets.randbelow(MAX_AES_KEY - MIN_AES_KEY) + MIN_AES_KEY).encode()

    def post(self, context, path, params, data):
        headers = {"Content-Type": "application/json"}
        cookies = {}
        if self.cookie:
            name, _, value = self.cookie.partition("=")
            if name and value:
                cookies[name] = value
        try:
            response = self.session.post(
                f"{self.host}{path}",
                params=params,
                data=data,
                headers=headers,
                cookies=cookies,
                verify=self.verify_ssl,
                timeout=self.timeout,
            )
        except RequestException:
            response = self.curl_post(path, params, data, cookies)
        if response.status_code == 403:
            self.clear_auth()
            raise RuntimeError(f"{context}: forbidden")
        response.raise_for_status()

        for cookie in response.headers.get("Set-Cookie", "").split(","):
            match = re.search(r"(sysauth=[A-Za-z0-9]+)", cookie)
            if match:
                self.cookie = match.group(1)
                break

        result = response.json()
        error_code = result.get("error_code")
        if error_code not in (None, "", 0):
            raise RuntimeError(f"{context}: response error_code={error_code}")
        return result

    def curl_post(self, path, params, data, cookies):
        url = f"{self.host}{path}"
        if params:
            query = "&".join(f"{key}={quote_plus(str(value))}" for key, value in params.items())
            url = f"{url}?{query}"
        command = [
            "curl",
            "-s",
            "-i",
            "--connect-timeout",
            str(self.timeout),
            "-X",
            "POST",
            url,
            "-H",
            "Content-Type: application/json",
            "--data-raw",
            data,
        ]
        if not self.verify_ssl:
            command.insert(2, "-k")
        if cookies:
            command.extend(["-H", "Cookie: " + "; ".join(f"{k}={v}" for k, v in cookies.items())])

        result = subprocess.run(command, check=True, capture_output=True, text=True, timeout=self.timeout + 5)
        head, separator, body = result.stdout.rpartition("\r\n\r\n")
        if not separator:
            head, _, body = result.stdout.rpartition("\n\n")
        status_match = re.search(r"HTTP/\S+\s+(\d+)", head)
        status = int(status_match.group(1)) if status_match else 200
        response = requests.Response()
        response.status_code = status
        response._content = body.encode()
        response.url = url
        for line in head.splitlines():
            if ":" not in line:
                continue
            key, value = line.split(":", 1)
            response.headers[key.strip()] = value.strip()
        return response

    def fetch_keys(self):
        response = self.post(
            "Fetch keys",
            "/cgi-bin/luci/;stok=/login",
            {"form": "keys"},
            json.dumps({"operation": "read"}),
        )
        keys = response["result"]["password"]
        self.password_rsa_n = int(keys[0], 16)
        self.password_rsa_e = int(keys[1], 16)

    def fetch_auth(self):
        response = self.post(
            "Fetch auth",
            "/cgi-bin/luci/;stok=/login",
            {"form": "auth"},
            json.dumps({"operation": "read"}),
        )
        auth = response["result"]
        self.sign_rsa_n = int(auth["key"][0], 16)
        self.sign_rsa_e = int(auth["key"][1], 16)
        self.seq = auth["seq"]

    def encode_payload(self, payload):
        payload_json = json.dumps(payload, separators=(",", ":")).encode()
        encrypted = AES.new(self.aes_key, AES.MODE_CBC, self.aes_iv).encrypt(
            pad(payload_json, AES.block_size)
        )
        data = base64.b64encode(encrypted).decode()
        sign = self.encode_sign(len(data))
        return f"sign={sign}&data={quote_plus(data)}"

    def encode_sign(self, data_len):
        auth_hash = hashlib.md5(f"{self.username}{self.password}".encode()).hexdigest()
        sign_text = (
            f"k={self.aes_key.decode()}&i={self.aes_iv.decode()}&h={auth_hash}&s={self.seq + data_len}"
        )
        return rsa_encrypt(self.sign_rsa_n, self.sign_rsa_e, sign_text.encode())

    def decrypt_data(self, context, data):
        if not data:
            self.clear_auth()
            raise RuntimeError(f"{context}: empty data")
        decrypted = AES.new(self.aes_key, AES.MODE_CBC, self.aes_iv).decrypt(
            base64.b64decode(data)
        )
        return json.loads(unpad(decrypted, AES.block_size).decode())

    def login(self):
        if self.aes_key is None:
            self.generate_aes_key_and_iv()
        if self.password_rsa_n is None:
            self.fetch_keys()
        if self.seq is None:
            self.fetch_auth()
        encrypted_password = rsa_encrypt(
            self.password_rsa_n,
            self.password_rsa_e,
            self.password.encode(),
        )
        response = self.post(
            "Login",
            "/cgi-bin/luci/;stok=/login",
            {"form": "login"},
            self.encode_payload({
                "operation": "login",
                "params": {"password": encrypted_password},
            }),
        )
        data = self.decrypt_data("Login", response["data"])
        if data.get("error_code") != 0:
            result = data.get("result") or {}
            attempts = result.get("attemptsAllowed", "unknown")
            raise RuntimeError(f"Login failed: error_code={data.get('error_code')}; attempts={attempts}")
        self.stok = data["result"]["stok"]
        if not self.cookie:
            raise RuntimeError("Login succeeded but no sysauth cookie was returned.")

    def clear_auth(self):
        self.seq = None
        self.stok = None
        self.cookie = None

    def call(self, context, section, form, payload):
        if not self.stok or not self.cookie:
            self.login()
        response = self.post(
            context,
            f"/cgi-bin/luci/;stok={self.stok}/admin/{section}",
            {"form": form},
            self.encode_payload(payload),
        )
        data = self.decrypt_data(context, response["data"])
        error_code = data.get("error_code") or data.get("errorcode")
        if error_code:
            raise RuntimeError(f"{context}: decoded error_code={error_code}")
        return data["result"]

    def list_decos(self):
        devices = self.call("List Devices", "device", "device_list", {"operation": "read"}).get("device_list", [])
        return [
            {
                "name": deco_name(device),
                "mac": device.get("mac"),
                "ip": device.get("device_ip"),
                "model": device.get("device_model"),
                "hardwareVersion": device.get("hardware_ver"),
                "firmwareVersion": device.get("software_ver"),
                "role": device.get("role"),
                "online": device.get("group_status") == "connected",
                "connectionType": device.get("connection_type"),
            }
            for device in devices
        ]

    def list_clients_for_deco(self, deco):
        clients = self.call(
            f"List Clients {deco['mac']}",
            "client",
            "client_list",
            {"operation": "read", "params": {"device_mac": deco["mac"]}},
        ).get("client_list", [])
        return [
            {
                "hostname": decode_name(client.get("name")),
                "mac": client.get("mac"),
                "ip": client.get("ip"),
                "connection": client.get("connection_type"),
                "interface": client.get("interface"),
                "upSpeed": client.get("up_speed"),
                "downSpeed": client.get("down_speed"),
                "active": client.get("online"),
                "linkedDecoMac": deco.get("mac"),
                "linkedDecoName": deco.get("name"),
                "linkedDecoRole": deco.get("role"),
            }
            for client in clients
            if client.get("online")
        ]


def run():
    if not config()["verify_ssl"]:
        requests.packages.urllib3.disable_warnings()
        ssl._create_default_https_context = ssl._create_unverified_context
    deco = DecoApi(**config())
    decos = deco.list_decos()
    clients = {}
    for node in decos:
        if not node.get("mac"):
            continue
        for client in deco.list_clients_for_deco(node):
            clients[client["mac"]] = client
    return {"decos": decos, "clients": list(clients.values())}


if __name__ == "__main__":
    try:
        print(json.dumps(run(), indent=2))
    except Exception as err:
        print(json.dumps({"error": str(err)}), file=sys.stderr)
        sys.exit(1)