# SPDX-FileCopyrightText: AISEC Pentesting Team
#
# SPDX-License-Identifier: Apache-2.0

import asyncio
import binascii
import io
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Protocol, Self
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse

from gallia.dumpcap import Dumpcap
from gallia.log import get_logger
from gallia.net import join_host_port
from gallia.transports.schemes import TransportScheme

logger = get_logger(__name__)


class TargetURI:
    """TargetURI represents a target to which gallia can connect.
    The target string must conform to a URI is specified by RFC3986.

    Basically, this is a wrapper around Python's ``urlparse()`` and
    ``parse_qs()`` methods. TargetURI provides frequently used properties
    for a more userfriendly usage. Instances are meant to be passed to
    :meth:`BaseTransport.connect()` of transport implementations.
    """

    def __init__(self, raw: str) -> None:
        self.raw = raw
        self.url = urlparse(raw)
        self.qs = parse_qs(self.url.query)

    @classmethod
    def from_parts(
        cls,
        scheme: str,
        host: str,
        port: int | None,
        args: dict[str, Any],
    ) -> Self:
        """Constructs a instance of TargetURI with the given arguments.
        The ``args`` dict is used for the query string.
        """
        netloc = host if port is None else join_host_port(host, port)
        return cls(urlunparse((scheme, netloc, "", "", urlencode(args), "")))

    @property
    def scheme(self) -> TransportScheme:
        """The URI scheme"""
        return TransportScheme(self.url.scheme)

    @property
    def hostname(self) -> str | None:
        """The hostname (without port)"""
        return self.url.hostname

    @property
    def port(self) -> int | None:
        """The port number"""
        return self.url.port

    @property
    def netloc(self) -> str:
        """The hostname and the portnumber, separated by a colon."""
        return self.url.netloc

    @property
    def path(self) -> str:
        """The path property of the url."""
        return self.url.path

    @property
    def location(self) -> str:
        """A URI string which only consists of the relevant scheme,
        the host and the port.
        """
        return f"{self.scheme}://{self.url.netloc}"

    @property
    def qs_flat(self) -> dict[str, str]:
        """A dict which contains the query string's key/value pairs.
        In case a key appears multiple times, this variant only
        contains the first found key/value pair. In contrast to
        :attr:`qs`, this variant avoids lists and might be easier
        to use for some cases.
        """
        d = {}
        for k, v in self.qs.items():
            d[k] = v[0]
        return d

    def __str__(self) -> str:
        return self.raw


class TransportProtocol(Protocol):
    mutex: asyncio.Lock
    target: TargetURI
    is_closed: bool

    def get_writer(self) -> asyncio.StreamWriter:
        raise NotImplementedError

    def get_reader(self) -> asyncio.StreamReader:
        raise NotImplementedError


class BaseTransport(ABC):
    """BaseTransport is the base class providing the required
    interface for all transports used by gallia.

    A transport usually is some kind of network protocol which
    carries an application level protocol. A good example is
    DoIP carrying UDS requests which acts as a minimal middleware
    on top of TCP.

    This class is to be used as a subclass with all abstractmethods
    implemented and the SCHEME property filled.

    A few methods provide a ``tags`` argument. The debug logs of these
    calls include these tags in the ``tags`` property of the relevant
    :class:`gallia.log.PenlogRecord`.
    """

    #: The scheme for the implemented protocol, e.g. "doip".
    SCHEME: str = ""
    #: The buffersize of the transport. Might be used in read() calls.
    #: Defaults to :const:`io.DEFAULT_BUFFER_SIZE`.
    BUFSIZE: int = io.DEFAULT_BUFFER_SIZE

    def __init__(self, target: TargetURI) -> None:
        self.check_scheme(target)
        self.mutex = asyncio.Lock()
        self.target = target
        self.is_closed = False
        self._dumpcap: Dumpcap | None = None

    def __init_subclass__(
        cls,
        /,
        scheme: str,
        bufsize: int = io.DEFAULT_BUFFER_SIZE,
        **kwargs: Any,
    ) -> None:
        super().__init_subclass__(**kwargs)
        cls.SCHEME = scheme
        cls.BUFSIZE = bufsize

    @classmethod
    def check_scheme(cls, target: TargetURI) -> None:
        """Checks if the provided URI has the correct scheme."""
        if target.scheme != cls.SCHEME:
            raise ValueError(f"invalid scheme: {target.scheme}; expected: {cls.SCHEME}")

    @abstractmethod
    async def connect(
        self,
        timeout: float | None = None,
    ) -> None:
        """Classmethod to connect the transport to a relevant target.
        The target argument is a URI, such as `doip://192.0.2.2:13400?src_addr=0xf4&dst_addr=0x1d"`
        An instance of the relevant transport class is returned.
        """

    @abstractmethod
    async def close(self) -> None:
        """Terminates the connection and clean up all allocated resources."""

    async def reconnect(self, timeout: float | None = None) -> None:
        """Closes the connection to the target and attempts to reconnect every
        100 ms until at max timeout. If timeout is None, only attempt to connect
        once.
        This method is safe for concurrent use.
        """
        async with self.mutex:
            try:
                await self.close()
            except ConnectionError as e:
                logger.warning(f"close() failed during reconnect ({e!r}); ignoring")

            async with asyncio.timeout(timeout):
                logger.debug(
                    f"Attempting to establish a new connection with a timeout of {timeout}"
                )
                while True:
                    try:
                        logger.debug("Attempting to connect...")
                        await self.connect()
                        return
                    except ConnectionError as e:
                        logger.info(f"Connection attempt failed while reconnecting: {e!r}")
                        if timeout is None:
                            logger.debug("Breaking out of the reconnect-loop since timeout is None")
                            raise e
                        await asyncio.sleep(0.1)

    @abstractmethod
    async def read(
        self,
        timeout: float | None = None,
        tags: list[str] | None = None,
    ) -> bytes:
        """Reads one message and returns its raw byte representation.
        An example for one message is 'one line, terminated by newline'
        for a TCP transport yielding lines.
        """

    @abstractmethod
    async def write(
        self,
        data: bytes,
        timeout: float | None = None,
        tags: list[str] | None = None,
    ) -> int:
        """Writes one message and return the number of written bytes."""

    async def request(
        self,
        data: bytes,
        timeout: float | None = None,
        tags: list[str] | None = None,
    ) -> bytes:
        """Chains a :meth:`write()` call with a :meth:`read()` call.
        The call is protected by a mutex and is thus safe for concurrent
        use.
        """
        async with self.mutex:
            return await self.request_unsafe(data, timeout, tags)

    async def request_unsafe(
        self,
        data: bytes,
        timeout: float | None = None,
        tags: list[str] | None = None,
    ) -> bytes:
        """Chains a :meth:`write()` call with a :meth:`read()` call.
        The call is **not** protected by a mutex. Only use this method
        when you know what you are doing.
        """
        await self.write(data, timeout, tags)
        return await self.read(timeout, tags)

    async def dumpcap_start(self, save_directory: Path) -> None:
        logger.debug("Attempting to start Dumpcap")

        if self._dumpcap is not None:
            logger.warning("Attempted to start Dumpcap, but it is already running!")
            return

        if (argument_list := await self.dumpcap_argument_list()) is None:
            logger.warning(
                f"Transport '{self.SCHEME}' of {self.target} does not support Dumpcap captures!"
            )
            return

        self._dumpcap = await Dumpcap.start(argument_list, save_directory)
        if self._dumpcap is None:
            logger.error("Dumpcap could not be started!")
        else:
            await self._dumpcap.sync()

    async def dumpcap_stop(self) -> None:
        if self._dumpcap is None:
            logger.debug("Attempting to stop Dumpcap, but it is not running!")
            return

        logger.info("Stopping Dumpcap")
        await self._dumpcap.stop()

    @abstractmethod
    async def dumpcap_argument_list(self) -> list[str] | None:
        """Return the arguments required for Dumpcap to capture traffic of this transport or `None` if not supported."""


class LinesTransportMixin:
    async def write(
        self: TransportProtocol,
        data: bytes,
        timeout: float | None = None,
        tags: list[str] | None = None,
    ) -> int:
        t = tags + ["write"] if tags is not None else ["write"]

        logger.trace(data.hex() + "0a", extra={"tags": t})

        writer = self.get_writer()
        writer.write(binascii.hexlify(data) + b"\n")
        await asyncio.wait_for(writer.drain(), timeout)
        return len(data)

    async def read(
        self: TransportProtocol,
        timeout: float | None = None,
        tags: list[str] | None = None,
    ) -> bytes:
        data = await asyncio.wait_for(self.get_reader().readline(), timeout)
        d = data.decode().strip()

        t = tags + ["read"] if tags is not None else ["read"]
        logger.trace(d + "0a", extra={"tags": t})

        return binascii.unhexlify(d)
