#!/usr/bin/env python3
"""Restore RB4011 identity and check WireGuard handshake after provision test."""
from __future__ import annotations

import time
import paramiko

JUMP = ("100.80.255.113", "otantik", "DeararbsFr2@")
MK = ("172.30.0.1", "admin", "demopass")


def wait_prompt(shell, timeout=20) -> str:
    buf = b""
    end = time.time() + timeout
    while time.time() < end:
        time.sleep(0.15)
        if shell.recv_ready():
            buf += shell.recv(65535)
            end = time.time() + 2.5
            if b"> " in buf[-200:]:
                break
    return buf.decode("utf-8", "replace")


def cmd(shell, line: str, timeout=20) -> str:
    shell.send(line + "\r")
    out = wait_prompt(shell, timeout)
    print("\n===== " + line + " =====")
    print(out[-3500:])
    return out


def main() -> None:
    jump = paramiko.SSHClient()
    jump.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    jump.connect(JUMP[0], username=JUMP[1], password=JUMP[2], timeout=25, allow_agent=False, look_for_keys=False)
    chan = jump.get_transport().open_channel("direct-tcpip", (MK[0], 22), ("127.0.0.1", 0))
    mk = paramiko.SSHClient()
    mk.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    mk.connect(MK[0], username=MK[1], password=MK[2], sock=chan, timeout=25, allow_agent=False, look_for_keys=False)
    shell = mk.invoke_shell(width=240, height=80)
    wait_prompt(shell, 12)

    cmd(shell, "/system identity set name=Otantik_Hub_RB4011")
    cmd(shell, "/system identity print")
    cmd(shell, "/interface wireguard print detail")
    cmd(shell, "/interface wireguard peers print detail where interface=wg-otantik")
    cmd(shell, "/ping 10.88.0.1 count=4")
    cmd(shell, "/ip hotspot walled-garden print where comment=wifi-otantik")

    mk.close()
    jump.close()


if __name__ == "__main__":
    main()