"""tcp.py

Created by Thomas Mangin on 2013-07-13.
Copyright (c) 2013-2017 Exa Networks. All rights reserved.
License: 3-clause BSD. (See the COPYRIGHT file)
"""

from __future__ import annotations

import re
import base64
import socket
import select
import platform

from struct import pack, calcsize
from typing import Iterator

from exabgp.util.errstr import errstr

from exabgp.protocol.family import AFI
from exabgp.protocol.ip import IP
from exabgp.reactor.network.error import errno
from exabgp.reactor.network.error import error

from exabgp.reactor.network.error import NotConnected
from exabgp.reactor.network.error import BindingError
from exabgp.reactor.network.error import MD5Error
from exabgp.reactor.network.error import NagleError
from exabgp.reactor.network.error import TTLError
from exabgp.reactor.network.error import AsyncError


def create(afi: AFI, interface: str | None = None) -> socket.socket:
    try:
        if afi == AFI.ipv4:
            io = socket.socket(socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP)
        if afi == AFI.ipv6:
            io = socket.socket(socket.AF_INET6, socket.SOCK_STREAM, socket.IPPROTO_TCP)
        try:
            io.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        except (OSError, AttributeError):
            pass
        try:
            io.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)  # pylint: disable=E1101
        except (OSError, AttributeError):
            pass

        if interface is not None:
            try:
                SO_BINDTODEVICE = getattr(socket, 'SO_BINDTODEVICE', 25)
                io.setsockopt(socket.SOL_SOCKET, SO_BINDTODEVICE, str(interface + '\0').encode('utf-8'))
            except OSError:
                raise NotConnected(f'Could not bind to device {interface}') from None
    except OSError:
        raise NotConnected('Could not create socket') from None
    return io


def bind(io: socket.socket, ip: str, afi: AFI) -> None:
    try:
        if afi == AFI.ipv4:
            io.bind((ip, 0))
        if afi == AFI.ipv6:
            io.bind((ip, 0, 0, 0))
    except OSError as exc:
        raise BindingError(f'Could not bind to local ip {ip} - {exc!s}') from None


def connect(io: socket.socket, ip: str, port: int, afi: AFI, md5: str) -> None:
    try:
        if afi == AFI.ipv4:
            io.connect((ip, port))
        if afi == AFI.ipv6:
            io.connect((ip, port, 0, 0))
    except OSError as exc:
        if exc.errno == errno.EINPROGRESS:
            return
        if md5:
            raise NotConnected(
                f'Could not connect to peer {ip}:{port}, check your MD5 password ({errstr(exc)})',
            ) from None
        raise NotConnected(f'Could not connect to peer {ip}:{port} ({errstr(exc)})') from None


# http://lxr.free-electrons.com/source/include/uapi/linux/tcp.h#L197
#
# #define TCP_MD5SIG_MAXKEYLEN    80
#
# struct tcp_md5sig {
# 	struct __kernel_sockaddr_storage tcpm_addr;     /* address associated */  128
# 	__u16   __tcpm_pad1;                            /* zero */                  2
# 	__u16   tcpm_keylen;                            /* key length */            2
# 	__u32   __tcpm_pad2;                            /* zero */                  4
# 	__u8    tcpm_key[TCP_MD5SIG_MAXKEYLEN];         /* key (binary) */         80
# }
#
# #define _K_SS_MAXSIZE   128
#
# #define _K_SS_ALIGNSIZE (__alignof__ (struct sockaddr *))
# /* Implementation specific desired alignment */
#
# typedef unsigned short __kernel_sa_family_t;
#
# struct __kernel_sockaddr_storage {
# 	__kernel_sa_family_t    ss_family;              /* address family */
# 	/* Following field(s) are implementation specific */
# 	char    __data[_K_SS_MAXSIZE - sizeof(unsigned short)];
# 	/* space to achieve desired size, */
# 	/* _SS_MAXSIZE value minus size of ss_family */
# } __attribute__ ((aligned(_K_SS_ALIGNSIZE)));   /* force desired alignment */


def md5(io: socket.socket, ip: str, port: int, md5: str, md5_base64: bool | None) -> None:
    platform_os = platform.system()
    if platform_os == 'FreeBSD':
        if md5:
            if md5 != 'kernel':
                raise MD5Error(
                    'FreeBSD requires that you set your MD5 key via ipsec.conf.\n'
                    'Something like:\n'
                    'flush;\n'
                    'add <local ip> <peer ip> tcp 0x1000 -A tcp-md5 "password";',
                )
            try:
                TCP_MD5SIG = 0x10
                io.setsockopt(socket.IPPROTO_TCP, TCP_MD5SIG, 1)
            except OSError:
                raise MD5Error(
                    'FreeBSD requires that you rebuild your kernel to enable TCP MD5 Signatures:\n'
                    'options         IPSEC\n'
                    'options         TCP_SIGNATURE\n'
                    'device          crypto\n',
                ) from None
    elif platform_os == 'Linux':
        try:
            md5_bytes = None
            if md5:
                if md5_base64 is True:
                    try:
                        md5_bytes = base64.b64decode(md5)
                    except TypeError:
                        raise MD5Error('Failed to decode base 64 encoded PSK') from None
                elif md5_base64 is None and not re.match('.*[^a-f0-9].*', md5):  # auto
                    options = [md5 + '==', md5 + '=', md5]
                    for md5 in options:
                        try:
                            md5_bytes = base64.b64decode(md5)
                            break
                        except TypeError:
                            pass

            # __kernel_sockaddr_storage
            n_af = IP.toaf(ip)
            n_addr = IP.pton(ip)
            n_port = socket.htons(port)

            # pack 'x' is padding, so we want the struct
            # Do not use '!' for the pack, the network (big) endian switch in
            # struct.pack is fighting against inet_pton and htons (note the n)

            if IP.toafi(ip) == AFI.ipv4:
                # SS_MAXSIZE is 128 but addr_family, port and ipaddr (8 bytes total) are written independently of the padding
                SS_MAXSIZE_PADDING = 128 - calcsize('HH4s')  # 8
                sockaddr = pack('HH4s%dx' % SS_MAXSIZE_PADDING, socket.AF_INET, n_port, n_addr)
            else:
                SS_MAXSIZE_PADDING = 128 - calcsize('HI16sI')  # 28
                SIN6_FLOWINFO = 0
                SIN6_SCOPE_ID = 0
                sockaddr = pack('HHI16sI%dx' % SS_MAXSIZE_PADDING, n_af, n_port, SIN6_FLOWINFO, n_addr, SIN6_SCOPE_ID)

            TCP_MD5SIG_MAXKEYLEN = 80
            TCP_MD5SIG = 14

            if md5_bytes:
                key = pack('2xH4x%ds' % TCP_MD5SIG_MAXKEYLEN, len(md5_bytes), md5_bytes)
                io.setsockopt(socket.IPPROTO_TCP, TCP_MD5SIG, sockaddr + key)
            elif md5:
                md5_bytes = bytes(md5, 'ascii')
                key = pack('2xH4x%ds' % TCP_MD5SIG_MAXKEYLEN, len(md5_bytes), md5_bytes)
                io.setsockopt(socket.IPPROTO_TCP, TCP_MD5SIG, sockaddr + key)
            # else:
            # 	key = pack('2xH4x%ds' % TCP_MD5SIG_MAXKEYLEN, 0, b'')
            # 	io.setsockopt(socket.IPPROTO_TCP, TCP_MD5SIG, sockaddr + key)

        except OSError as exc:
            if exc.errno != errno.ENOENT:
                raise MD5Error(
                    'This linux machine does not support TCP_MD5SIG, you can not use MD5 ({})'.format(errstr(exc))
                ) from None
    elif md5:
        raise MD5Error('ExaBGP has no MD5 support for {}'.format(platform_os))


# TCP-AO (RFC 5925) - TCP Authentication Option
# Linux kernel 6.7+ required
#
# From linux/tcp.h:
# struct tcp_ao_add {
#     struct __kernel_sockaddr_storage addr;  /* 128 bytes */
#     char alg_name[64];                       /* algorithm name */
#     __s32 ifindex;                           /* VRF interface */
#     __u32 set_current:1, set_rnext:1, reserved:30;
#     __u16 reserved2;
#     __u8 prefix, sndid, rcvid, maclen, keyflags, keylen;
#     __u8 key[TCP_AO_MAXKEYLEN];              /* 80 bytes */
# };

# Algorithm name mapping: user-friendly config -> kernel crypto name
TCP_AO_ALGORITHMS = {
    'hmac-sha-1-96': 'hmac(sha1)',
    'aes-128-cmac-96': 'cmac(aes)',
    'hmac-sha-256': 'hmac(sha256)',
}


def tcp_ao(
    io: socket.socket,
    ip: str,
    port: int,
    password: str,
    keyid: int,
    algorithm: str,
    base64_encoded: bool = False,
) -> None:
    """Set TCP-AO (RFC 5925) authentication on socket.

    Args:
        io: Socket to configure
        ip: Remote IP address
        port: Remote port (0 for listening sockets)
        password: Authentication key
        keyid: Key ID (0-255, used for both sndid and rcvid)
        algorithm: Algorithm name (hmac-sha-1-96, aes-128-cmac-96, hmac-sha-256)
        base64_encoded: Whether password is base64 encoded

    Raises:
        TCPAOError: If TCP-AO cannot be configured
    """
    from exabgp.reactor.network.error import TCPAOError

    if not password:
        return

    platform_os = platform.system()
    if platform_os != 'Linux':
        raise TCPAOError(f'TCP-AO is only supported on Linux (current: {platform_os})')

    # Validate keyid
    if not (0 <= keyid <= 255):
        raise TCPAOError(f'TCP-AO keyid must be 0-255 (got: {keyid})')

    # Map algorithm name
    kernel_alg = TCP_AO_ALGORITHMS.get(algorithm)
    if kernel_alg is None:
        valid = ', '.join(TCP_AO_ALGORITHMS.keys())
        raise TCPAOError(f'Invalid TCP-AO algorithm "{algorithm}". Valid: {valid}')

    # Decode password if base64 encoded
    if base64_encoded:
        try:
            key_bytes = base64.b64decode(password)
        except (TypeError, ValueError) as e:
            raise TCPAOError(f'Failed to decode base64 TCP-AO key: {e}') from None
    else:
        key_bytes = password.encode('utf-8')

    # Validate key length
    TCP_AO_MAXKEYLEN = 80
    if len(key_bytes) > TCP_AO_MAXKEYLEN:
        raise TCPAOError(f'TCP-AO key too long ({len(key_bytes)} > {TCP_AO_MAXKEYLEN} bytes)')

    try:
        # Socket option constants (from linux/tcp.h)
        TCP_AO_ADD_KEY = 38
        ALG_NAME_LEN = 64

        # Build sockaddr_storage (same as MD5)
        n_af = IP.toaf(ip)
        n_addr = IP.pton(ip)
        n_port = socket.htons(port)

        if IP.toafi(ip) == AFI.ipv4:
            SS_MAXSIZE_PADDING = 128 - calcsize('HH4s')  # 120
            sockaddr = pack('HH4s%dx' % SS_MAXSIZE_PADDING, socket.AF_INET, n_port, n_addr)
        else:
            SS_MAXSIZE_PADDING = 128 - calcsize('HHI16sI')  # 100
            SIN6_FLOWINFO = 0
            SIN6_SCOPE_ID = 0
            sockaddr = pack('HHI16sI%dx' % SS_MAXSIZE_PADDING, n_af, n_port, SIN6_FLOWINFO, n_addr, SIN6_SCOPE_ID)

        # Pack algorithm name (64 bytes, null-padded)
        alg_name = kernel_alg.encode('ascii').ljust(ALG_NAME_LEN, b'\x00')

        # Pack remaining fields:
        # ifindex (4 bytes) + flags (4 bytes) + reserved2 (2 bytes)
        # + prefix, sndid, rcvid, maclen, keyflags, keylen (6 bytes)
        # + key (80 bytes)
        ifindex = 0  # No VRF
        flags = 0x3  # set_current=1, set_rnext=1
        reserved2 = 0
        prefix = 0  # Exact match (no prefix filtering)
        sndid = keyid
        rcvid = keyid
        maclen = 0  # Use default for algorithm
        keyflags = 0
        keylen = len(key_bytes)

        # Pad key to TCP_AO_MAXKEYLEN
        key_padded = key_bytes.ljust(TCP_AO_MAXKEYLEN, b'\x00')

        # Pack: ifindex(i) + flags(I) + reserved2(H) + prefix,sndid,rcvid,maclen,keyflags,keylen(6B) + key(80s)
        ao_data = pack(
            'iIH6B%ds' % TCP_AO_MAXKEYLEN,
            ifindex,
            flags,
            reserved2,
            prefix,
            sndid,
            rcvid,
            maclen,
            keyflags,
            keylen,
            key_padded,
        )

        # Combine sockaddr + alg_name + ao_data
        tcp_ao_add = sockaddr + alg_name + ao_data

        io.setsockopt(socket.IPPROTO_TCP, TCP_AO_ADD_KEY, tcp_ao_add)

    except OSError as exc:
        if exc.errno == errno.ENOPROTOOPT:
            raise TCPAOError(
                f'This Linux kernel does not support TCP-AO (requires kernel 6.7+). Error: {errstr(exc)}'
            ) from None
        raise TCPAOError(f'Failed to set TCP-AO option: {errstr(exc)}') from None


def nagle(io: socket.socket, ip: str) -> None:
    try:
        # diable Nagle's algorithm (no grouping of packets)
        io.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    except (OSError, AttributeError):
        raise NagleError("Could not disable nagle's algorithm for {}".format(ip)) from None


def ttl(io: socket.socket, ip: str, ttl: int | None) -> None:
    # None (ttl-security unset) or zero (maximum TTL) is the same thing
    if ttl:
        try:
            io.setsockopt(socket.IPPROTO_IP, socket.IP_TTL, ttl)
        except OSError as exc:
            raise TTLError(f'This OS does not support IP_TTL (ttl-security) for {ip} ({errstr(exc)})') from None


def ttlv6(io: socket.socket, ip: str, ttl: int | None) -> None:
    if ttl:
        try:
            io.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_UNICAST_HOPS, ttl)
        except OSError as exc:
            raise TTLError(
                'This OS does not support unicast_hops (ttl-security) for {} ({})'.format(ip, errstr(exc))
            ) from None


def min_ttl(io: socket.socket, ip: str, ttl: int | None) -> None:
    # None (ttl-security unset) or zero (maximum TTL) is the same thing
    if ttl:
        try:
            IP_MINTTL = getattr(socket, 'IP_MINTTL', None)
            if IP_MINTTL is not None:
                io.setsockopt(socket.IPPROTO_IP, IP_MINTTL, ttl)
            else:
                raise AttributeError('IP_MINTTL not available')
        except OSError as exc:
            raise TTLError(
                'This OS does not support IP_MINTTL (ttl-security) for {} ({})'.format(ip, errstr(exc))
            ) from None
        except AttributeError:
            pass

        try:
            io.setsockopt(socket.IPPROTO_IP, socket.IP_TTL, ttl)
        except OSError as exc:
            raise TTLError(
                'This OS does not support IP_MINTTL or IP_TTL (ttl-security) for {} ({})'.format(ip, errstr(exc)),
            ) from None


def min_ttlv6(io: socket.socket, ip: str, ttl: int | None) -> None:
    # None (ttl-security unset) or zero (maximum TTL) is the same thing
    if ttl:
        try:
            # IPV6_MINHOPCOUNT is Linux-specific (value 73), not always in the socket module
            IPV6_MINHOPCOUNT = getattr(socket, 'IPV6_MINHOPCOUNT', 73)
            io.setsockopt(socket.IPPROTO_IPV6, IPV6_MINHOPCOUNT, ttl)
        except OSError as exc:
            raise TTLError(
                'This OS does not support IPV6_MINHOPCOUNT (ttl-security) for {} ({})'.format(ip, errstr(exc))
            ) from None

        try:
            io.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_UNICAST_HOPS, ttl)
        except OSError as exc:
            raise TTLError(
                'This OS does not support IPV6_UNICAST_HOPS (ttl-security) for {} ({})'.format(ip, errstr(exc)),
            ) from None


def asynchronous(io: socket.socket, ip: str) -> None:
    try:
        io.setblocking(False)
    except OSError as exc:
        raise AsyncError('could not set socket non-blocking for {} ({})'.format(ip, errstr(exc))) from None


def ready(io: socket.socket) -> Iterator[tuple[bool, str]]:
    poller = select.poll()
    poller.register(io, select.POLLOUT | select.POLLNVAL | select.POLLERR)

    found = False

    while True:
        try:
            for _, event in poller.poll(0):
                if event & select.POLLOUT or event & select.POLLIN:
                    found = True
                elif event & select.POLLHUP:
                    yield False, 'could not connect, retrying'
                    return
                elif event & select.POLLERR or event & select.POLLNVAL:
                    yield False, 'connect attempt failed, issue with reading on the network, retrying'
                    return

            if found:
                err = io.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
                if not err:
                    yield True, 'connection established'
                    return
                elif err in error.block:
                    yield False, 'connect attempt failed, retrying, reason {}'.format(errno.errorcode[err])
                    return
            yield False, 'waiting for socket to become ready'
        except OSError as err:
            yield False, 'error, retrying {}'.format(str(err))
            return
