#!/usr/bin/env python3

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

"""The client for the backdoor API of the website."""

from __future__ import annotations

import ast
import asyncio
import http.client
import io
import os
import pickle  # nosec: B403
import pydoc
import re  # pylint: disable=preferred-module
import socket
import sys
import time
import traceback
import uuid
from base64 import b64encode
from collections.abc import Callable, Iterable, MutableMapping
from contextlib import suppress
from importlib import import_module
from textwrap import dedent
from types import EllipsisType
from typing import Any, Required, TypeAlias, TypedDict, cast
from urllib.parse import SplitResult, quote, quote_plus, urlsplit

with suppress(ModuleNotFoundError):
    # pylint: disable=shadowed-import
    import dill as pickle  # type: ignore[import-untyped, no-redef]  # noqa: F811, B950  # nosec: B403

try:
    import hy  # type: ignore[import-untyped]
except ModuleNotFoundError:
    hy = None  # pylint: disable=invalid-name

try:
    import idna
except ModuleNotFoundError:
    idna = None  # type: ignore[assignment]  # pylint: disable=invalid-name

try:
    import socks  # type: ignore[import-untyped]
except ModuleNotFoundError:
    socks = None  # pylint: disable=invalid-name

if os.environ.get("DISABLE_UVLOOP") not in {"y", "yes", "t", "true", "on", "1"}:
    with suppress(ModuleNotFoundError):
        asyncio.set_event_loop_policy(import_module("uvloop").EventLoopPolicy())

E = eval(  # pylint: disable=eval-used  # nosec: B307
    "eval(repr((_:=[],_.append(_))[0]))[0][0]"
)

ErrorTuple: TypeAlias = tuple[int, str]

FLUFL = True


class Response(TypedDict):  # noqa: D101
    # pylint: disable=missing-class-docstring
    success: bool | EllipsisType
    output: None | str
    result: None | tuple[str, None | bytes] | SystemExit


class Proxy(TypedDict, total=False):  # noqa: D101
    # pylint: disable=missing-class-docstring
    type: Required[int]
    addr: Required[str]
    port: int
    rdns: bool
    username: str
    password: str


async def create_socket(
    addr: str,
    port: int | str,
    proxy: None | Proxy,
) -> socket.socket:
    """Create a socket (optionally with a proxy)."""
    # pylint: disable=too-complex
    if proxy is not None and not socks:
        raise NotImplementedError("PySocks is required for proxy support")
    loop = asyncio.get_running_loop()
    address_infos = await loop.getaddrinfo(
        addr,
        port,
        # PySocks doesn't support AF_INET6
        family=0 if proxy is None else socket.AF_INET,
        type=socket.SOCK_STREAM,
    )
    if not address_infos:
        raise OSError("getaddrinfo() returned empty list")
    exceptions = []
    for address_info in address_infos:
        sock = None
        try:
            sock = cast(
                socket.socket,
                (
                    socks.socksocket(address_info[0])
                    if proxy is not None
                    else socket.socket(address_info[0])
                ),
            )
            sock.setblocking(False)
            if proxy is not None:
                sock.set_proxy(  # type: ignore[attr-defined]
                    proxy["type"],
                    proxy["addr"],
                    proxy.get("port"),
                    proxy.get("rdns", True),
                    proxy.get("username"),
                    proxy.get("password"),
                )
            await loop.sock_connect(sock, address_info[4])
            return sock
        except OSError as exc:
            if sock is not None:
                sock.close()
            exceptions.append(exc)
            continue
        except BaseException:
            if sock is not None:
                sock.close()
            raise
    if len(exceptions) == 1:
        raise exceptions[0]
    model = str(exceptions[0])
    if all(str(exc) == model for exc in exceptions):
        raise exceptions[0]
    raise OSError(
        f"Multiple exceptions: {', '.join(str(exc) for exc in exceptions)}"
    )


async def request(  # noqa: C901
    method: str,
    url: str | SplitResult,
    headers: None | MutableMapping[str, str] = None,
    body: None | bytes | bytearray | Iterable[bytes | bytearray] | str = None,
    proxy: None | Proxy = None,
) -> tuple[int, dict[str, str], bytes]:
    """Insanely awesome HTTP client."""
    # pylint: disable=line-too-long, too-complex
    # pylint: disable=too-many-branches, too-many-locals, while-used
    if isinstance(url, str):
        url = urlsplit(url)
    if url.scheme not in {"", "http", "https"}:
        raise ValueError(f"Unsupported scheme: {url.scheme}")
    if not url.hostname:
        raise ValueError("URL has no hostname")
    if headers is None:
        headers = {}
    if isinstance(body, str):
        body = body.encode("UTF-8")
    if isinstance(body, memoryview):
        body = body.tobytes()
    if isinstance(body, Iterable) and not isinstance(body, (bytes, bytearray)):
        body = b"".join(body)  # type: ignore[arg-type]
    https = url.scheme == "https"
    header_names = [x.strip().title() for x in headers.keys()]
    if "Host" not in header_names:
        host: None | str = None
        if idna:  # type: ignore[truthy-bool]
            with suppress(idna.core.InvalidCodepoint):
                host = idna.encode(url.hostname).decode("ASCII")
                host = f"{host}:{url.port}" if url.port else host
        if not host:
            host = url.netloc.encode("IDNA").decode("ASCII")
        headers["Host"] = host
    if body and "Content-Length" not in header_names:
        headers["Content-Length"] = str(len(body))
    sock = await create_socket(
        url.hostname,
        url.port or ("https" if https else "http"),
        proxy,
    )
    reader, writer = await asyncio.open_connection(
        sock=sock,
        ssl=https,
        server_hostname=url.hostname if https else None,
    )
    writer.write(
        (
            method
            + " "
            + (quote(url.path) or "/")
            + ("?" + quote_plus(url.query) if url.query else "")
            + " HTTP/1.0\r\n"
        ).encode("ASCII")
        + "\r\n".join(
            [f"{key}:{value}" for key, value in headers.items()] + [""]
        ).encode("LATIN-1")
        + b"\r\n"
    )
    if body:
        writer.write(body)
    await writer.drain()
    del headers
    e, data = E, b""
    while chunk := await reader.read():
        if b"\r\n\r\n" in (data := data + chunk) and e is E:
            e, data = data.split(b"\r\n\r\n", 1)
            status, o = re.match(r"HTTP/.+? (\d+).*?\r\n(.*)", e.decode("LATIN-1"), 24).groups()  # type: ignore[union-attr]  # noqa: B950
            headers = dict(re.match(r"([^\s]+):\s*(.+?)\s*$", x, 24).groups() for x in o.split("\r\n"))  # type: ignore[union-attr, misc]  # noqa: B950
    writer.close()
    await writer.wait_closed()
    if "status" not in locals():
        raise AssertionError("No HTTP response received")
    # pylint: disable-next=possibly-used-before-assignment
    return int(status), headers, data  # type: ignore[possibly-undefined]


def detect_mode(code: str) -> str:
    """Detect which mode needs to be used."""
    import __future__  # pylint: disable=import-outside-toplevel

    flags: int = ast.PyCF_ONLY_AST
    if FLUFL:
        flags |= __future__.barry_as_FLUFL.compiler_flag

    try:
        compile(code, "", "eval", flags, cast(bool, 0x5F3759DF), 0)
        return "eval"
    except SyntaxError:
        compile(code, "", "exec", flags, cast(bool, 0x5F3759DF), 0)
        return "exec"


def send(
    url: str | SplitResult,
    key: str,
    code: str,
    mode: str = "exec",
    session: None | str = None,
    proxy: None | Proxy = None,
) -> tuple[int, dict[str, str], Response | ErrorTuple | str | None | bytes]:
    """Send code to the backdoor API."""
    # pylint: disable=too-many-arguments
    body = code.encode("UTF-8")
    if isinstance(url, str):
        url = urlsplit(url)
    if not url.path:
        url = url._replace(path="/api/backdoor")
    key = f"Bearer {b64encode(key.encode('UTF-8')).decode('ASCII')}"
    headers = {
        "Authorization": key,
        "Accept": (
            "application/vnd.uqfoundation.dill"
            if pickle.__name__ == "dill"
            else "application/vnd.python.pickle"
        ),
        "X-Pickle-Protocol": str(pickle.HIGHEST_PROTOCOL),
        "X-Future-Feature": "annotations",
    }
    if FLUFL:
        headers["X-Future-Feature"] += ", barry_as_FLUFL"
    if session:
        headers["X-Backdoor-Session"] = session
    response = asyncio.run(
        request(
            "POST",
            url._replace(path=f"{url.path.removesuffix('/')}/{mode}"),
            headers,
            body,
            proxy,
        )
    )
    try:
        return (
            response[0],  # status
            response[1],  # header
            pickle.loads(response[2]),  # data  # nosec: B301
        )
    except pickle.UnpicklingError:
        return response


def lisp_always_active() -> bool:
    """Return True if LISP is always active."""
    return (
        hy
        and not hy.eval(
            hy.read(
                '(* (- (* (+ 0 1) 2 3 4 5) (+ 6 7 8 9 10 11)) '  # fmt: skip
                '(int (= (. (__import__ "os.path") sep) "/")))'
            )
        )
        and not int.from_bytes(
            getattr(
                os,
                "洀漀搀渀愀爀甀".encode("UTF-16-BE")[::-1].decode("UTF-16-BE"),
            )(1),
            "big",
        )
        // (69 // 4 - 1)
    )


def run_and_print(  # noqa: C901
    url: str,
    key: str,
    code: str,
    lisp: bool = False,
    session: None | str = None,
    proxy: None | Proxy = None,
    time_requests: bool = False,
    *,
    # pylint: disable=redefined-builtin
    print: Callable[..., None] = print,
) -> None:
    """Run the code and print the output."""
    # pylint: disable=too-complex, too-many-arguments, too-many-branches
    # pylint: disable=too-many-locals, too-many-statements
    start_time = time.monotonic()
    if lisp or lisp_always_active():
        code = hy.disassemble(hy.read(code), True)
    try:
        response = send(
            url,
            key,
            code,
            detect_mode(code),
            session,
            proxy,
        )
    except SyntaxError as exc:
        print("".join(traceback.format_exception_only(exc)).strip())
        return
    if time_requests:
        took = time.monotonic() - start_time
        if took > 1:
            color = "91"  # red
        elif took > 0.5:
            color = "93"  # yellow
        else:
            color = "92"  # green
        print(f"\033[{color}mTook: {took:.3f}s\033[0m")
    status, headers, body = response  # pylint: disable=unused-variable
    if status >= 400:
        reason = (
            body[1]
            if isinstance(body, tuple)
            else http.client.responses[status]
        )
        print("\033[91m" + f"{status} {reason}" + "\033[0m")
        if isinstance(body, str):
            print("\033[91m" + body + "\033[0m")
    elif body is None:  # pylint: disable=confusing-consecutive-elif
        pass
    elif isinstance(body, dict):
        if isinstance(body["success"], bool):
            print(f"Success: {body['success']}")
        if isinstance(body["output"], str) and body["output"]:
            print("Output:")
            print(body["output"].strip())
        if isinstance(body["result"], SystemExit):
            raise body["result"]
        if isinstance(body["result"], tuple):
            if not body["success"]:
                print(body["result"][0])
                return
            result_obj: Any = None
            if isinstance(body["result"][1], bytes):
                try:
                    result_obj = pickle.loads(body["result"][1])  # nosec: B301
                except Exception:  # pylint: disable=broad-except
                    if sys.flags.dev_mode:
                        traceback.print_exc()
            if (
                isinstance(result_obj, tuple)
                and len(result_obj) == 2
                and result_obj[0] == "PagerTuple"
                and isinstance(result_obj[1], str)
            ):
                pydoc.pager(result_obj[1])
            else:
                print("Result:")
                print(body["result"][0])
    else:
        print(f"Response has unexpected type {type(body).__name__}!")
        print(body)


def shellify(code: str) -> str:
    """Modify code in a way that it gets executed in a shell."""
    if not code.startswith("!"):
        return code

    return f"""
async def run_shell_50821273052022fbc283():
    from os.path import dirname
    import asyncio, sys, os
    path = ":".join(filter(bool, [dirname(sys.executable), os.environ.get("PATH")]))
    proc = await asyncio.create_subprocess_shell(
        r'''{code[1:].strip().replace("'''", """''' "'''" r'''""")} '''.strip(),
        asyncio.subprocess.DEVNULL,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
        env={{"PATH": path}},
    )
    output, _ = await asyncio.wait_for(proc.communicate(), 60 * 60)
    if output:
        print(output.decode("UTF-8"))
await run_shell_50821273052022fbc283()
del run_shell_50821273052022fbc283
"""


def main() -> int | str:  # noqa: C901
    """Parse arguments, load the config and start the backdoor client."""
    # pylint: disable=too-complex, too-many-branches
    # pylint: disable=too-many-locals, too-many-statements

    # pylint: disable-next=import-outside-toplevel, import-error, useless-suppression
    from pyrepl.python_reader import (  # type: ignore[import, unused-ignore]
        ReaderConsole,
        main as _main,
    )

    if "--help" in sys.argv or "-h" in sys.argv:
        print(
            dedent(
                """\
                Accepted arguments:

                    --dev              use a separate config for a local dev instance
                    --lisp             enable Lots of Irritating Superfluous Parentheses
                    --new-proxy        don't use the saved proxy
                    --new-session      start a new session with saved URL and key
                    --no-config        start without loading/saving the config
                    --no-patch-help    don't patch help()
                    --reset-config     reset the whole config
                    --timing           print the time it took to execute each command

                    --help or -h       show this help message
                """
            ),
            end="",
            file=sys.stderr,
        )
        return 0
    for arg in sys.argv[1:]:
        if arg not in {
            "--dev",
            "--lisp",
            "--new-proxy",
            "--new-session",
            "--no-config",
            "--no-patch-help",
            "--reset-config",
            "--timing",
            "--help",
            "-h",
        }:
            print(f"Unknown argument: {arg}")
            sys.exit(64 + 4 + 1)
    file: io.IOBase
    url: None | str = None
    key: None | str = None
    session: None | str = None
    proxy_type: None | int | EllipsisType = None
    proxy_addr: None | str = None
    proxy_port: None | int = None
    proxy_rdns: None | bool = True
    proxy_username: None | str = None
    proxy_password: None | str = None
    config_pickle = os.path.join(
        os.path.expanduser(os.getenv("XDG_CONFIG_HOME") or "~/.config"),
        "an-backdoor-client/"
        + ("dev-" if "--dev" in sys.argv else "")
        + "session.pickle",
    )
    if "--reset-config" in sys.argv and os.path.exists(config_pickle):
        os.remove(config_pickle)
    if "--no-config" not in sys.argv:
        try:
            with open(config_pickle, "rb") as file:
                config = pickle.load(file)  # nosec: B301
        except FileNotFoundError:
            pass
        else:
            url = config.get("url")
            key = config.get("key")
            session = config.get("session")
            proxy_type = config.get("proxy_type")
            proxy_addr = config.get("proxy_addr")
            proxy_port = config.get("proxy_port")
            proxy_rdns = config.get("proxy_rdns")
            proxy_username = config.get("proxy_username")
            proxy_password = config.get("proxy_password")
            if "--new-session" in sys.argv:
                print(f"Using URL {url}")
            else:
                print(f"Using URL {url} with existing session {session}")
    while not url:  # pylint: disable=while-used
        url = input("URL: ").strip().rstrip("/")
        if not url:
            print("No URL given!")
        elif "://" not in url:
            if not url.startswith("//"):
                url = "//" + url
            if re.match(
                r"^(\/\/)(localhost|127\.0\.0\.1|\[::1\])(\:\d+)?(/\S*)?$", url
            ):
                url = "http:" + url
            else:
                url = "https:" + url
            print(f"Using URL {url}")

    while not key:  # pylint: disable=while-used
        key = input("Key: ").strip()
        if not key:
            print("No key given!")

    if proxy_type is None or "--new-proxy" in sys.argv:
        proxy_url_str = input("Proxy (leave empty for none): ").strip()
        if proxy_url_str:
            if "://" not in proxy_url_str:
                if not proxy_url_str.startswith("//"):
                    proxy_url_str = "//" + proxy_url_str
                proxy_url_str = "socks5:" + proxy_url_str
            proxy_url = urlsplit(proxy_url_str)
            if proxy_url.hostname:
                proxy_type = int(socks.PROXY_TYPES[proxy_url.scheme.upper()])
                proxy_addr = proxy_url.hostname
                proxy_port = proxy_url.port
                proxy_rdns = True
                proxy_username = proxy_url.username or None
                proxy_password = proxy_url.password or None
            else:
                print("Invalid proxy URL!")
        else:
            print("No proxy given!")
    if isinstance(proxy_type, EllipsisType):
        proxy_type = None
    if proxy_type is not None:
        print(
            f"Using {socks.PRINTABLE_PROXY_TYPES[proxy_type]} proxy "
            f"{proxy_addr}{f':{proxy_port}' if proxy_port else ''}"
            + (f" with username {proxy_username}" if proxy_username else "")
        )
    else:
        print("Using no proxy (use --new-proxy to be able to set one)")

    if not session or "--new-session" in sys.argv:
        session = input("Session (enter nothing for random session): ")
        if not session:
            session = str(uuid.uuid4())
        print(f"Using session {session}")

    if "--no-config" not in sys.argv:
        os.makedirs(os.path.dirname(config_pickle), exist_ok=True)
        with open(config_pickle, "wb") as file:
            pickle.dump(
                {
                    "url": url,
                    "key": key,
                    "session": session,
                    "proxy_type": proxy_type or ...,  # not None (None == ask)
                    "proxy_addr": proxy_addr,
                    "proxy_port": proxy_port,
                    "proxy_rdns": proxy_rdns,
                    "proxy_username": proxy_username,
                    "proxy_password": proxy_password,
                },
                file,
                5,
            )

    proxy: None | Proxy

    if proxy_type and proxy_addr:
        proxy = {"type": proxy_type, "addr": proxy_addr}
        if proxy_port:
            proxy["port"] = proxy_port
        if proxy_rdns:
            proxy["rdns"] = proxy_rdns
        if proxy_username and proxy_password:
            proxy["username"] = proxy_username
            proxy["password"] = proxy_password
    else:
        proxy = None

    def send_to_remote(code: str, *, mode: str) -> Any:
        """Send code to the remote backdoor and return the unpickled body."""
        return send(
            url,
            key,
            code,
            mode,
            session,
            proxy,
        )[2]

    if "--no-patch-help" not in sys.argv:
        body = send_to_remote(
            # fmt: off
            "class _HelpHelper_92005ecf3788faea8346a7919fba0232188561ab:\n"
            "  def __call__(self, *args, **kwargs):\n"
            "    import io\n"
            "    import pydoc\n"
            "    helper_output = io.StringIO()\n"
            "    pydoc.Helper(io.StringIO(), helper_output)(*args, **kwargs)\n"
            "    return 'PagerTuple', helper_output.getvalue()\n"
           f"  __str__ = __repr__ = lambda _:{repr(help)!r}\n"  # noqa: E131
            "help = _HelpHelper_92005ecf3788faea8346a7919fba0232188561ab()\n"
            "del _HelpHelper_92005ecf3788faea8346a7919fba0232188561ab",
            # fmt: on
            mode="exec",
        )
        if not (isinstance(body, dict) and body["success"]):
            print("\033[91mPatching help() failed!\033[0m")

    if "--lisp" in sys.argv:
        if not hy:
            sys.exit("\033[91mHy is not installed!\033[0m")
        body = send_to_remote("__import__('hy')", mode="exec")
        if not (isinstance(body, dict) and body["success"]):
            print("\033[91mInjecting Hy builtins failed!\033[0m")

    body = send_to_remote(
        "import sys\nprint('Python', sys.version, 'on', sys.platform)",
        mode="exec",
    )
    if isinstance(body, dict) and body["success"] and body["output"]:
        print(f"\033[92mConnection to {url} was successful.\033[0m")
        print(body["output"].strip())
    else:
        print("\033[91mGetting remote information failed.\033[0m")
    print(
        'Type "copyright", "credits" or '
        'use the "help" function for more information.'
    )

    def _run_and_print(  # type: ignore[no-any-unimported]
        self: ReaderConsole,
        code: str,
    ) -> None:
        # pylint: disable=unused-argument
        try:
            run_and_print(
                url,
                key,
                shellify(code),
                "--lisp" in sys.argv,
                session,
                proxy,
                "--timing" in sys.argv,
            )
        except Exception:  # pylint: disable=broad-except
            print(
                "\033[91mAn unexpected error occurred. "
                "Please contact a developer.\033[0m"
            )
            traceback.print_exc()

    # patch the reader console to use our function
    rc_execute = ReaderConsole.execute
    try:
        ReaderConsole.execute = _run_and_print
        # run the reader
        _main(print_banner=False, clear_main=False)
    except EOFError:
        pass
    finally:
        # restore the original method
        ReaderConsole.execute = rc_execute

    return 0


if __name__ == "__main__":
    sys.exit(main())
