import os
import socket
import typing
import ctypes
import contextlib
import struct
import binascii
import enum
import re
import time
import base64
from ctypes import c_char, c_short
from typing import Tuple
import asyncio
from aioupnp.upnp import UPnP, UPnPError, get_gateway_and_lan_addresses
from aioupnp.constants import SSDP_IP_ADDRESS
import miniupnpc


_IFF_PROMISC  = 0x0100
_SIOCGIFFLAGS = 0x8913  # get the active flags
_SIOCSIFFLAGS = 0x8914  # set the active flags
_ETH_P_ALL    = 0x0003  # all protocols
ETHER_HEADER_LEN = 6 + 6 + 2
VLAN_HEADER_LEN = 2

printable = re.compile(b"([a-z0-9!\"#$%&'()*+,.\/:;<=>?@\[\] ^_`{|}~-]*)")


class PacketTypes(enum.Enum):  # if_packet.h
    HOST = 0
    BROADCAST = 1
    MULTICAST = 2
    OTHERHOST = 3
    OUTGOING = 4
    LOOPBACK = 5
    FASTROUTE = 6


class Layer2(enum.Enum):  # https://www.iana.org/assignments/ieee-802-numbers/ieee-802-numbers.xhtml
    IPv4 = 0x0800
    ARP =  0x0806
    VLAN = 0x8100
    MVRP = 0x88f5
    MMRP = 0x88f6
    IPv6 = 0x86dd
    GRE  = 0xb7ea


class Layer3(enum.Enum):  # https://www.iana.org/assignments/protocol-numbers/protocol-numbers.xhtml
    ICMP = 1
    IGMP = 2
    TCP = 6
    UDP = 17


class _ifreq(ctypes.Structure):
    _fields_ = [("ifr_ifrn", c_char * 16),
                ("ifr_flags", c_short)]


@contextlib.contextmanager
def _promiscuous_posix_socket_context(interface: str):
    import fcntl  # posix-only
    sock = socket.socket(socket.PF_PACKET, socket.SOCK_RAW, socket.htons(_ETH_P_ALL))
    ifr = _ifreq()
    ifr.ifr_ifrn = interface.encode()[:16]
    fcntl.ioctl(sock, _SIOCGIFFLAGS, ifr)  # get the flags
    ifr.ifr_flags |= _IFF_PROMISC  # add the promiscuous flag
    fcntl.ioctl(sock, _SIOCSIFFLAGS, ifr)  # update
    sock.setblocking(False)
    try:
        yield sock
    finally:
        ifr.ifr_flags ^= _IFF_PROMISC  # mask it off (remove)
        fcntl.ioctl(sock, _SIOCSIFFLAGS, ifr)  # update
        print("closed posix promiscuous socket")


@contextlib.contextmanager
def _promiscuous_non_posix_socket_context():
    # the public network interface
    HOST = socket.gethostbyname(socket.gethostname())
    # create a raw socket and bind it to the public interface
    sock = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.IPPROTO_IP)
    # prevent socket from being left in TIME_WAIT state, enabling reuse
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind((HOST, 0))
    # Include IP headers
    sock.setsockopt(socket.IPPROTO_IP, socket.IP_HDRINCL, 1)
    # receive all packages
    sock.ioctl(socket.SIO_RCVALL, socket.RCVALL_ON)
    sock.setblocking(False)
    try:
        yield sock
    finally:
        # disable promiscuous mode
        sock.ioctl(socket.SIO_RCVALL, socket.RCVALL_OFF)
        print("closed non-posix promiscuous socket")


def promiscuous(interface: typing.Optional[str] = None) -> typing.ContextManager[socket.socket]:
    if os.name == 'posix':
        return _promiscuous_posix_socket_context(interface)
    return _promiscuous_non_posix_socket_context()


def ipv4_to_str(addr: bytes) -> str:
    return ".".join((str(b) for b in addr))


def pretty_mac(mac: bytes) -> str:
    return ":".join((('0' if b < 16 else '') + hex(b)[2:] for b in mac))


def split_byte(b: int, bit=4) -> Tuple[bytes, bytes]:
    return chr(((b >> (8-bit)) % 256) << (8-bit) >> (8-bit)).encode(), chr(((b << bit) % 256) >> bit).encode()


class EtherFrame:
    __slots__ = [
        'source_mac',
        'target_mac',
        'ether_type',
        'vlan_id',
        'tpid'
    ]

    def __init__(self, source_mac: bytes, target_mac: bytes, ether_type: int,
                 vlan_id: typing.Optional[int] = None, tpid: typing.Optional[int] = None):
        self.source_mac = source_mac
        self.target_mac = target_mac
        self.ether_type = ether_type
        self.vlan_id = vlan_id
        self.tpid = tpid

    def encode(self) -> bytes:
        if self.vlan_id is None:
            return struct.pack("6s6sH", *(getattr(self, slot) for slot in self.__slots__[:-2]))
        return struct.pack("6s6sHHH", *(getattr(self, slot) for slot in self.__slots__))

    @classmethod
    def decode(cls, packet: bytes) -> Tuple['EtherFrame', bytes]:
        vlan_id = None
        tpid = None
        if struct.unpack(f'!H', packet[12:14])[0] == Layer2.VLAN.value:
            target_mac, source_mac, tpid, vlan_id, ether_type, data = struct.unpack(f'!6s6sHHH{len(packet) - ETHER_HEADER_LEN - VLAN_HEADER_LEN}s', packet)
        else:
            target_mac, source_mac, ether_type, data = struct.unpack(f'!6s6sH{len(packet) - ETHER_HEADER_LEN}s', packet)
        return cls(source_mac, target_mac, ether_type, vlan_id, tpid), data

    def debug(self) -> str:
        if self.vlan_id is None:
            return f"EtherFrame(source={pretty_mac(self.source_mac)}, target={pretty_mac(self.target_mac)}, " \
                f"ether_type={Layer2(self.ether_type).name})"
        return f"EtherFrame(source={pretty_mac(self.source_mac)}, target={pretty_mac(self.target_mac)}, " \
               f"ether_type={Layer2(self.ether_type).name}, vlan={self.vlan_id})"


class IPv4Packet:
    __slots__ = [
        'ether_frame',
        'version',
        'header_length',
        'dscp',
        'ecn',
        'total_length',
        'identification',
        'df',
        'mf',
        'flag',
        'fragment_offset',
        'ttl',
        'protocol',
        'header_checksum',
        '_source_address',
        '_destination_address',
        'data',
        'packet_type',
        'interface'
    ]

    ETHER_TYPE = Layer2.IPv4

    def __init__(self, ether_frame: EtherFrame, version: int, header_length: int, dscp: int, ecn: int,
                 total_length: int, identification: int, mf: bool, df: bool, flag: bool, fragment_offset: int, ttl: int,
                 protocol: int, header_checksum: int, source_address: bytes, destination_address: bytes, data: bytes,
                 packet_type: int, interface: str):
        self.ether_frame = ether_frame
        self.version = version
        self.header_length = header_length
        self.dscp = dscp
        self.ecn = ecn
        self.total_length = total_length
        self.identification = identification
        self.mf = mf
        self.df = df
        self.flag = flag
        self.fragment_offset = fragment_offset
        self.ttl = ttl
        self.protocol = Layer3(protocol)
        self.header_checksum = header_checksum
        self._source_address = source_address
        self._destination_address = destination_address
        self.data = data
        self.packet_type = PacketTypes(packet_type)
        self.interface = interface

    @property
    def source(self) -> str:
        return ipv4_to_str(self._source_address)

    @property
    def destination(self) -> str:
        return ipv4_to_str(self._destination_address)

    @staticmethod
    def checksum(header: bytes) -> int:
        c = 0
        for i in range(0, len(header), 2):
            c += int.from_bytes(header[i:i + 2], 'big')
            while c > 0xffff:
                c %= 0xffff
        while c > 0xffff:
            c %= 0xffff
        return c ^ 0xffff

    def get_header(self) -> bytes:
        version_and_hlen = (self.version << 4) + self.header_length
        dscp_and_ecn = (self.dscp << 2) + self.ecn
        flags = (4 if self.flag else 0) + (2 if self.df else 0) + (1 if self.mf else 0)
        df_mf_and_fragment = (flags << 12) + self.fragment_offset
        return struct.pack(
            '!BBHHHBBH4s4s', version_and_hlen, dscp_and_ecn, self.total_length,
            self.identification, df_mf_and_fragment, self.ttl, self.protocol.value,
            self.header_checksum, self._source_address, self._destination_address
        )

    @classmethod
    def decode(cls, ether_frame: EtherFrame, packet: bytes, packet_type: int, interface: str) -> 'IPv4Packet':
        if cls.checksum(packet[:20]):
            raise ValueError(f'\nipv4 checksum failed, frame: {ether_frame.debug()}\n'
                             f'packet: {binascii.hexlify(packet).decode()}, checksum: {hex(cls.checksum(packet[:20]))}')
        data_len = len(packet) - 20
        version_and_hlen, dscp_and_ecn, tlen, ident, df_mf_and_fragment, ttl, proto, checksum, source, dest = \
            struct.unpack(
                f'!BBHHHBBH4s4s', packet[:20]
            )
        version, hlen = split_byte(version_and_hlen)
        flags = df_mf_and_fragment >> 13
        mask = (flags << 13) | df_mf_and_fragment
        fragment = mask ^ df_mf_and_fragment
        flag, df, mf = False, False, False
        if flags % 2:
            mf = True
            flags -= 1
        if flags % 2:
            df = True
            flags -= 2
        if flags % 4:
            flag = True
            flags -= 4
        dscp, ecn = split_byte(dscp_and_ecn, 6)
        return cls(
            ether_frame, ord(version), ord(hlen), ord(dscp), ord(ecn), tlen, ident, mf, df, flag, fragment, ttl,
            proto, checksum, source, dest, packet[20:], packet_type, interface
        )

    def encode(self) -> bytes:
        return self.ether_frame.encode() + self.get_header() + self.data

    @property
    def printable_data(self) -> str:
        return b".".join(printable.findall(self.data)).decode()

    def __repr__(self) -> str:
        return f"IPv4(protocol={self.protocol.name}, " \
               f"iface={self.interface}, " \
               f"type={self.packet_type.name}, " \
               f"source={ipv4_to_str(self._source_address)}, " \
               f"destination={ipv4_to_str(self._destination_address)}, " \
               f"data_len={len(self.data)})"


def make_filter(l3_protocol=None, src=None, dst=None, invert=False):
    def filter_packet(packet: IPv4Packet):
        if l3_protocol and not Layer3(packet.protocol) == l3_protocol:
            return False
        if src and not packet.source == src:
            return False
        if dst and not packet.destination == dst:
            return False
        return True
    if invert:
        return lambda packet: not filter_packet(packet)
    return filter_packet


async def sniff_ipv4(filters=None, kill=None):
    start = time.perf_counter()
    loop = asyncio.get_event_loop()

    async def sock_recv(sock, n):
        """Receive data from the socket.

        The return value is a bytes object representing the data received.
        The maximum amount of data to be received at once is specified by
        nbytes.
        """
        if loop._debug and sock.gettimeout() != 0:
            raise ValueError("the socket must be non-blocking")
        fut = loop.create_future()
        _sock_recv(fut, None, sock, n)
        return await fut

    def _sock_recv(fut, registered_fd, sock, n):
        # _sock_recv() can add itself as an I/O callback if the operation can't
        # be done immediately. Don't use it directly, call sock_recv().
        if registered_fd is not None:
            # Remove the callback early.  It should be rare that the
            # selector says the fd is ready but the call still returns
            # EAGAIN, and I am willing to take a hit in that case in
            # order to simplify the common case.
            loop.remove_reader(registered_fd)
        if fut.cancelled():
            return
        try:
            data, flags = sock.recvfrom(n)
        except (BlockingIOError, InterruptedError):
            fd = sock.fileno()
            loop.add_reader(fd, _sock_recv, fut, fd, sock, n)
        except Exception as exc:
            fut.set_exception(exc)
        else:
            fut.set_result((data, flags))

    with promiscuous('lo') as sock:
        while True:
            if not kill:
                data, flags = await sock_recv(sock, 9000)
            else:
                t = asyncio.create_task(sock_recv(sock, 9000))
                await asyncio.wait([t, kill.wait()], return_when=asyncio.FIRST_COMPLETED)
                if kill.is_set():
                    break
                data, flags = await t
            # https://stackoverflow.com/questions/42821309/how-to-interpret-result-of-recvfrom-raw-socket/45215859#45215859
            if data:
                try:
                    ether_frame, packet = EtherFrame.decode(data)
                    if ether_frame.ether_type == Layer2.IPv4.value:
                        interface, _, packet_type, _, _ = flags
                        ipv4 = IPv4Packet.decode(ether_frame, packet, packet_type, interface)
                        if not filters or any((f(ipv4) for f in filters)):
                            yield time.perf_counter() - start, ipv4
                except ValueError:
                    pass


async def main():
    loop = asyncio.get_event_loop()
    gateway, lan = get_gateway_and_lan_addresses('default')

    done = asyncio.Event()

    def discover_aioupnp():
        async def _discover():
            print("test aioupnp")
            try:
                u = await UPnP.discover()
                try:
                    await u.get_external_ip()
                except UPnPError:
                    pass
                try:
                    await u.get_redirects()
                except UPnPError:
                    pass
                try:
                    external_port = await u.get_next_mapping(1234, 'TCP', 'aioupnp testing')
                except UPnPError:
                    external_port = None
                try:
                    await u.get_redirects()
                except UPnPError:
                    pass
                if external_port:
                    try:
                        await u.delete_port_mapping(external_port, 'TCP')
                    except UPnPError:
                        pass
                    try:
                        await u.get_redirects()
                    except UPnPError:
                        pass
            finally:
                print("done with aioupnp test")
        asyncio.create_task(_discover())

    def discover_miniupnpc():
        def _miniupnpc_discover():
            try:
                u = miniupnpc.UPnP()
            except:
                return
            try:
                u.discover()
            except:
                return
            try:
                u.selectigd()
            except:
                return
            try:
                u.externalipaddress()
            except:
                return

        async def _discover():
            print("test miniupnpc")
            try:
                await loop.run_in_executor(None, _miniupnpc_discover)
            finally:
                done.set()
                print("done with miniupnpc test")

        asyncio.create_task(_discover())

    loop.call_later(2, discover_aioupnp)
    loop.call_later(10, discover_miniupnpc)
    start = time.perf_counter()
    f = open("upnp_packet_cap.txt", "w")
    try:
        async for (ts, ipv4_packet) in sniff_ipv4([
                make_filter(l3_protocol=Layer3.UDP, src=SSDP_IP_ADDRESS),
                make_filter(l3_protocol=Layer3.UDP, dst=SSDP_IP_ADDRESS),
                make_filter(l3_protocol=Layer3.UDP, src=lan, dst=gateway),
                make_filter(l3_protocol=Layer3.UDP, src=gateway, dst=lan),
                make_filter(l3_protocol=Layer3.TCP, src=lan, dst=gateway),
                make_filter(l3_protocol=Layer3.TCP, src=gateway, dst=lan)], done):
            f.write(f"{time.perf_counter() - start},{ipv4_packet.packet_type.name},{ipv4_packet.source},{ipv4_packet.destination},{base64.b64encode(ipv4_packet.data).decode()}\n")
            print(ts, ipv4_packet)
            print(ipv4_packet.printable_data)
            print()
    except KeyboardInterrupt:
        print("stopping")
    finally:
        f.close()


if __name__ == "__main__":
    asyncio.run(main())