#!/usr/bin/env python3

import argparse
import base64
import json
import logging
import ssl
import time
import urllib.error
import urllib.request

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
logger = logging.getLogger(__name__)

ROUTER_HOST = "192.168.1.1"
ROUTER_USERNAME = "abdus"
ROUTER_PASSWORD = "xAsametk50"
PUBLIC_IP_URL = "https://api.ipify.org"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Drop and re-establish the router's PPPoE session to request a new CGNAT public IP, without rebooting",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--host", default=ROUTER_HOST)
    parser.add_argument("--timeout", type=float, default=15.0, help="Router request timeout in seconds")
    parser.add_argument("--wait-after-disconnect", type=float, default=5.0, help="Seconds to wait before reconnecting")
    parser.add_argument("--max-wait", type=float, default=60.0, help="Seconds to wait for the session to come back up")
    parser.add_argument("--skip-public-ip-check", action="store_true", help="Do not call an external service to compare public IPs")
    parser.add_argument("--dry-run", action="store_true", help="Log in and find the active PPPoE session, but do not disconnect it")
    return parser.parse_args()


class RouterSession:
    """Logs in to the router on enter and always logs out on exit, so a crash never leaks a session.

    The router has a limit on concurrent logins and returns "Maxium number of login account has
    reached" once it's hit; the only way out of that is waiting for idle sessions to time out.
    """

    def __init__(self, host: str, timeout: float) -> None:
        self.base_url = f"https://{host}"
        self.timeout = timeout
        self.session_key = ""
        # The router uses a self-signed certificate. Over HTTPS the UI skips its RSA/AES login layer.
        context = ssl.create_default_context()
        context.check_hostname = False
        context.verify_mode = ssl.CERT_NONE
        self.opener = urllib.request.build_opener(
            urllib.request.HTTPSHandler(context=context),
            urllib.request.HTTPCookieProcessor(),
        )

    def request(self, path: str, body: dict | None, method: str) -> dict:
        separator = "&" if "?" in path else "?"
        suffix = f"{separator}sessionkey={self.session_key}" if self.session_key else ""
        headers = {"Content-Type": "application/json"}
        if self.session_key and method != "GET":
            # Since the ABSB.5 firmware, state-changing requests (POST/PUT/DELETE) are only
            # accepted with the session key repeated in this header; the ?sessionkey= query
            # param alone (still required for GETs) gets a bare 403 on everything else.
            headers["CSRFToken"] = self.session_key
        request = urllib.request.Request(
            url=f"{self.base_url}{path}{suffix}",
            data=json.dumps(body).encode() if body is not None else None,
            headers=headers,
            method=method,
        )
        with self.opener.open(request, timeout=self.timeout) as response:
            raw = response.read().decode()
        return json.loads(raw) if raw.strip() else {}

    def login(self) -> None:
        body = self.request(
            path="/UserLogin",
            body={
                "Input_Account": ROUTER_USERNAME,
                "Input_Passwd": base64.b64encode(ROUTER_PASSWORD.encode()).decode(),
                "currLang": "en",
                "RememberPassword": 0,
                "SHA512_password": False,
            },
            method="POST",
        )
        if body.get("result") != "ZCFG_SUCCESS":
            raise RuntimeError(f"Login rejected: {body}")
        self.session_key = body.get("sessionkey", "")
        logger.info("Logged in as %s", ROUTER_USERNAME)

    def logout(self) -> None:
        if not self.session_key:
            return
        try:
            self.request(path="/cgi-bin/UserLogout", body={}, method="POST")
            logger.info("Logged out, session released")
        except Exception as exc:
            logger.warning("Logout failed, the router may hold this session until it times out: %s", exc)
        finally:
            self.session_key = ""

    def get_status(self) -> dict:
        return self.request(path="/cgi-bin/DAL?oid=status", body=None, method="GET")

    def wan_action(self, wanppp_idx: int, action: str) -> None:
        # Disconnect/Connect tear down or rebuild the PPPoE link, so the router often doesn't get
        # an HTTP response back in time (or at all, if the link it was answering over just dropped).
        # That doesn't mean the action failed: treat a timeout as "sent, outcome unknown" and let
        # the caller confirm the real state by polling, instead of raising and giving up.
        try:
            reply = self.request(path="/cgi-bin/DAL?oid=status", body={"action": action, "wanpppIdx": wanppp_idx}, method="PUT")
        except (TimeoutError, urllib.error.URLError, ConnectionError) as exc:
            logger.info("%s request did not get a response (%s), checking the actual state instead", action, exc)
            return
        if reply.get("result") not in (None, "ZCFG_SUCCESS"):
            raise RuntimeError(f"Router rejected {action}: {reply}")

    def __enter__(self) -> "RouterSession":
        self.login()
        return self

    def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None:
        self.logout()


def find_active_pppoe(status: dict) -> dict:
    candidates = [
        entry
        for entry in status["Object"][0]["WanLanInfo"]
        if entry.get("Encapsulation") == "PPPoE" and entry.get("Status") == "Up"
    ]
    if not candidates:
        raise RuntimeError("No active PPPoE WAN session found to refresh")
    if len(candidates) > 1:
        logger.warning("Multiple active PPPoE sessions found, using the first: %s", candidates[0].get("Name"))
    return candidates[0]


def wan_ip_of(entry: dict) -> str | None:
    return next(iter(entry.get("IPv4Address") or []), {}).get("IPAddress")


def find_wan_entry(status: dict, wanppp_idx: int) -> dict | None:
    return next(
        (e for e in status["Object"][0]["WanLanInfo"] if e.get("wanpppIdx") == wanppp_idx and e.get("Encapsulation") == "PPPoE"),
        None,
    )


def poll_until_up(session: RouterSession, wanppp_idx: int, timeout: float) -> dict | None:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            entry = find_wan_entry(status=session.get_status(), wanppp_idx=wanppp_idx)
            if entry and entry.get("Status") == "Up" and wan_ip_of(entry=entry):
                return entry
        except (urllib.error.URLError, json.JSONDecodeError, ConnectionError, KeyError, IndexError) as exc:
            logger.info("Still reconnecting: %s", exc)
        time.sleep(2.0)
    return None


def ensure_reconnected(session: RouterSession, wanppp_idx: int, max_wait: float, connect_retries: int = 2) -> dict:
    """Poll for the link to come back; if it's stuck at N/A (the router's own auto-reconnect can
    fail to kick in), nudge it with another Connect call instead of leaving it down."""
    per_attempt = max(max_wait / (connect_retries + 1), 15.0)
    for attempt in range(connect_retries + 1):
        entry = poll_until_up(session=session, wanppp_idx=wanppp_idx, timeout=per_attempt)
        if entry:
            return entry
        if attempt < connect_retries:
            logger.warning("Not back up yet, retrying Connect (%d/%d)...", attempt + 1, connect_retries)
            session.wan_action(wanppp_idx=wanppp_idx, action="Connect")
    raise TimeoutError(f"WAN session did not come back up within {max_wait:.0f}s, even after retrying Connect")


def fetch_public_ip(timeout: float) -> str | None:
    try:
        with urllib.request.urlopen(PUBLIC_IP_URL, timeout=timeout) as response:
            return response.read().decode().strip()
    except Exception as exc:
        logger.info("Could not fetch public IP from %s: %s", PUBLIC_IP_URL, exc)
        return None


def main() -> None:
    args = parse_args()
    public_ip_before = None if args.skip_public_ip_check else fetch_public_ip(timeout=args.timeout)
    if public_ip_before:
        logger.info("Public IP before: %s", public_ip_before)

    with RouterSession(host=args.host, timeout=args.timeout) as session:
        wan = find_active_pppoe(status=session.get_status())
        wanppp_idx = wan["wanpppIdx"]
        logger.info("Active PPPoE session: %s (wanpppIdx=%s, current WAN IP=%s)", wan.get("Name"), wanppp_idx, wan_ip_of(entry=wan))

        if args.dry_run:
            logger.info("Dry run, not disconnecting")
            return

        logger.info("Disconnecting PPPoE session...")
        session.wan_action(wanppp_idx=wanppp_idx, action="Disconnect")

        time.sleep(args.wait_after_disconnect)

        logger.info("Reconnecting PPPoE session...")
        session.wan_action(wanppp_idx=wanppp_idx, action="Connect")

        entry = ensure_reconnected(session=session, wanppp_idx=wanppp_idx, max_wait=args.max_wait)
        logger.info("Reconnected. WAN IP: %s", wan_ip_of(entry=entry))

        if public_ip_before:
            public_ip_after = fetch_public_ip(timeout=args.timeout)
            if public_ip_after:
                logger.info("Public IP after: %s", public_ip_after)
                if public_ip_after == public_ip_before:
                    logger.warning("Public IP did not change (ISP returned the same CGNAT address)")
                else:
                    logger.info("Public IP changed: %s -> %s", public_ip_before, public_ip_after)


if __name__ == "__main__":
    main()