from typing import Tuple
from HCI_tools import *

import serial
import socket as skt
from binascii import hexlify
from math import ceil
import struct as s
import copy, time
import ctypes
import os
from dataclasses import dataclass

HCI_DLL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'HCI_parser', 'HCI_parser.dll')

def cmd_seq(l):
    return [hex_read(x) for x in l]

def adv_cmds(random, adv_interval_ms = 50):
    int_bytes = hp(s.pack('<H', int(round(1.6 * adv_interval_ms))))
    random_byte = hp(s.pack('B', 1 if random else 0))
    return cmd_seq([
        '01 06 20 0F ' + int_bytes + ' ' + int_bytes + ' 00 ' + random_byte + ' 00 00 00 00 00 00 00 07 00', # LE Set Advertising Parameters
        # '01 08 20 20 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00', #LE Set Advertising Data
        '01 0A 20 01 01', # LE Set Advertising Enable
        ])

def connect_cmd(random, addr):
    addr_bytes = hp(s.pack('>Q', addr))[6:]
    random_byte = hp(s.pack('B', 1 if random else 0))
    return cmd_seq([
        '01 0D 20 19 E0 01 E0 01 00 ' + random_byte + ' ' + addr_bytes + ' ' + random_byte + ' F0 00 F0 00 00 00 E8 03 00 00 00 00', # LE Create Connection
        ])

def set_random_addr(addr):
    addr_bytes = hp(s.pack('>Q', addr))[6:]
    return cmd_seq(['01 05 20 06 ' + addr_bytes])

def add_to_filter_accept(addr):
    addr_bytes = hp(s.pack('>Q', addr))[6:]
    return cmd_seq(['01 11 20 07 00 ' + addr_bytes])

def disconnect_cmd(handle):
    return cmd_seq(['01 06 04 03 ' + hp(s.pack('<H', handle)) + ' 13'])

def read_BD_ADDR_cmd():
    return cmd_seq(['01 09 10 00'])

def init_cmds():
    return cmd_seq([
        '01 03 0C 00', # Reset
        '01 09 10 00', # Read BD_ADDR
        '01 01 0C 08 FF FF FF FF FF FF FF 3F', # BB Set Event Mask
        '01 01 20 08 FF FF FF FF FF FF FF FF', # LE Set Event Mask
        '01 74 20 02 2F 01', # Set Host Feature (for CS)
        ])

def encrypt_cmd(handle, LTK):
    key_bytes = hp(s.pack('<Q', LTK & (2**64 - 1))) + ' ' + hp(s.pack('<Q', LTK >> 64))
    return cmd_seq([
        '01 19 20 1C ' + hp(s.pack('<H', handle)) + ' 00 00 00 00 00 00 00 00 00 00 ' + key_bytes, # LE Enable Encryption
        ])

def respond_to_connection_cmds(handle, key):
    handle_bytes = hp(s.pack('<H', handle))
    key_bytes = hp(s.pack('<Q', key & (2**64 - 1))) + ' ' + hp(s.pack('<Q', key >> 64))
    return cmd_seq([
        # 04 3E 0D 05 80 00 LE LTK Key Request
        '01 1A 20 12 ' + handle_bytes + ' ' + key_bytes, # LE LTK Request Reply
        ])

def HCI_send(com, payl_out, print_reply=True):
    com.write(payl_out)
    # ret = com.read_until()
    ret = com.read(1000)
    if print_reply:
        print(hexlify(ret))
    else:
        return ret
    
@dataclass
class testerConfig:
    com_IDs: Tuple[int]

    adv_addr: int = 0x010203040506
    cen_addr: int = 0x0708090A0B0C

    pcapng_port: int | None = None # 25000
    pcap_ports: Tuple[int] | None = (25005, 25006)
    btsnoop_ports: Tuple[int] | None = None # (25001, 25002)
    raw_ports: Tuple[int] | None = None # (25003, 25004)

    HCI_decode: bool = True
    HCI_verbose: bool = False
    print_in: bool = True
    print_out: bool = False
    print_logs: bool = True
    HCI_to_pcapng: bool = False
    random_addresses : bool = True

    btsnoop_fname: str | None = None

    bauds: Tuple[int] = (115200, 115200)
    timeout: float = 0.05
    LTK: int = 5


def config_no_sockets(com_IDs):
    return testerConfig(com_IDs=com_IDs, pcapng_port=None, btsnoop_ports=None, raw_ports=None, pcap_ports = None)

class devTester:
    def __init__(self, c : testerConfig):
        self.com_IDs = c.com_IDs
        self.coms = []
        self.inputs = []
        self.input_msgs = []
        for id, baud in zip(c.com_IDs, c.bauds):
            self.coms.append(serial.Serial('COM{}'.format(id), baud, timeout=c.timeout))
            self.inputs.append([])
            self.input_msgs.append([])

        self.btsnoop_fname = c.btsnoop_fname
        self.print_out = c.print_out
        self.print_in = c.print_in
        self.input_cbs = []
        self.HCI_decode = c.HCI_decode
        self.HCI_verbose = c.HCI_verbose
        self.print_logs = c.print_logs
        self.random_addresses = c.random_addresses
        
        self.cen_addr = c.cen_addr
        self.adv_addr = c.adv_addr
        self.conn_handles = [None, None]
        self.encrypted = False
        self.encrypted_count = 0
        self.LTK = c.LTK
        self.HCI_to_pcapng = c.HCI_to_pcapng

        if self.HCI_decode:
            self.HCI_parser = ctypes.cdll.LoadLibrary(HCI_DLL_PATH)
            self.HCI_parser.HCI_summary_string.restype = ctypes.c_char_p
            self.HCI_parser.HCI_tree_string.restype = ctypes.c_char_p

        self.pcapng_port_num = c.pcapng_port
        self.pcap_port_nums = c.pcap_ports
        self.btsnoop_port_nums = c.btsnoop_ports
        self.raw_port_nums = c.raw_ports

        if self.pcapng_port_num is not None:
            self.pcapng_port = skt.socket(skt.AF_INET, skt.SOCK_STREAM)
            try:
                self.pcapng_port.connect(("127.0.0.1", self.pcapng_port_num))
                self.pcapng_interfaces = (0, 1)
                self.pcapng_port.sendall(pcapng_interface_description())
                self.pcapng_port.sendall(pcapng_interface_description())
            except ConnectionRefusedError as e:
                print("Failed to connect to pcapng port: {}".format(e))
                self.pcapng_port = None
        else:
            self.pcapng_port = None

        if self.pcap_port_nums is not None:
            self.pcap_ports = []
            try:
                for i in range(2):
                    self.pcap_ports.append(skt.socket(skt.AF_INET, skt.SOCK_STREAM))
                    self.pcap_ports[i].connect(("127.0.0.1", self.pcap_port_nums[i]))
                    self.pcap_ports[i].sendall(pcap_file_header())
            except ConnectionRefusedError as e:
                print("Failed to connect to pcap ports: {}".format(e))
                self.pcap_ports = None
        else:
            self.pcap_ports = None

        if self.btsnoop_port_nums is not None:
            try:
                self.btsnoop_ports = []
                for i in range(2):
                    self.btsnoop_ports.append(skt.socket(skt.AF_INET, skt.SOCK_STREAM))
                    self.btsnoop_ports[i].connect(("127.0.0.1", self.btsnoop_port_nums[i]))
                    self.btsnoop_ports[i].sendall(b'btsnoop\x00')
                    self.btsnoop_ports[i].sendall(s.pack('>II', 1, 1002))    
            except ConnectionRefusedError as e:
                print("Failed to connect to btsnoop ports: {}".format(e))
                self.btsnoop_ports = None                    
        else:
            self.btsnoop_ports = None

        if self.raw_port_nums is not None:
            self.raw_ports = []
            try:
                for i in range(2):
                    self.raw_ports.append(skt.socket(skt.AF_INET, skt.SOCK_STREAM))
                    self.raw_ports[i].connect(("127.0.0.1", self.raw_port_nums[i]))
            except ConnectionRefusedError as e:
                print("Failed to connect to raw ports: {}".format(e))
                self.raw_ports = None                  
                
        else:
            self.raw_ports = None

        if c.btsnoop_fname is not None:
            self.btsnoop = open(c.btsnoop_fname, 'wb')
            start_btsnoop(self.btsnoop)
        else:
            self.btsnoop = None

        for idx in range(len(self.coms)):
            self.read_version(idx)

        self.read_all()

    # Override this if your device needs extra initialisation
    def send_init_cmds(self):
        self.send_parallel_cmd_seqs((init_cmds(), init_cmds()))

    def handle_input(self, com_idx, in_bytes):
        if self.raw_ports is not None:
            self.raw_ports[com_idx].sendall(in_bytes)

        self.inputs[com_idx].extend(in_bytes)
        while True:
            try:
                N = HCI_pkt_len(self.inputs[com_idx])
            except:
                return

            if N is None or len(self.inputs[com_idx]) < N:
                return

            pkt = HCI_pkt(copy.deepcopy(self.inputs[com_idx][:N]))
            self.inputs[com_idx] = self.inputs[com_idx][N:]

            
            us_since_epoch = int(round(time.time() * 1000000))

            self.input_msgs[com_idx].append(pkt)
            if self.btsnoop:
                write_btsnoop_pkt(self.btsnoop, False, pkt.bytes, us_since_epoch)

            if self.btsnoop_ports is not None:
                self.btsnoop_ports[com_idx].sendall(build_btsnoop_pkt(False, pkt.bytes, us_since_epoch))

            if self.pcap_ports is not None:
                pcap_pkt = build_pcap_pkt(False, pkt.bytes, us_since_epoch)
                # print('Write pcap packet {}'.format(hp(pcap_pkt)))
                self.pcap_ports[com_idx].sendall(pcap_pkt)

            if self.pcapng_port is not None and self.HCI_to_pcapng:
                self.pcapng_port.sendall(pcapng_h4_phdr(self.pcapng_interfaces[com_idx], False, pkt.bytes, us_since_epoch))

            if self.print_in:
                if self.HCI_decode:
                    if self.HCI_verbose:
                        tree = self.HCI_parser.HCI_tree_string(bytes(pkt.bytes), len(pkt.bytes))
                        print('    <-{}: {}'.format(self.coms[com_idx].port, tree.decode('utf-8')))
                    else:
                        summ = self.HCI_parser.HCI_summary_string(bytes(pkt.bytes), len(pkt.bytes), True)
                        print('    <-{}: {}'.format(self.coms[com_idx].port, truncate(summ.decode('utf-8'), 250)))
                else:
                    print('<- {: 4} == {} bytes on {}: {}'.format(len(pkt.bytes), N, self.coms[com_idx].port, hp(pkt.bytes)))

            for cb in self.input_cbs:
                cb(com_idx, pkt)


    def send_cmd(self, com_idx, payl):
        if self.raw_ports is not None:
            self.raw_ports[com_idx].sendall(bytes(payl))
        
        self.coms[com_idx].write(payl)

        us_since_epoch = int(round(time.time() * 1000000))
        
        if self.print_out:
            print('-> ' + hp(payl))
            if self.HCI_decode:
                if self.HCI_verbose:
                    tree = self.HCI_parser.HCI_tree_string(bytes(payl), len(payl))
                    print('    ->{}: {}'.format(self.coms[com_idx].port, tree.decode('utf-8')))
                else:
                    summ = self.HCI_parser.HCI_summary_string(bytes(payl), len(payl), True)
                    print('    ->{}: {}'.format(self.coms[com_idx].port, truncate(summ.decode('utf-8'), 250)))

        if self.btsnoop_ports is not None:
            self.btsnoop_ports[com_idx].sendall(build_btsnoop_pkt(True, payl, us_since_epoch))            

        if self.btsnoop:
            write_btsnoop_pkt(self.btsnoop, True, payl, us_since_epoch)        

        if self.pcap_ports is not None:
            self.pcap_ports[com_idx].sendall(build_pcap_pkt(True, payl, us_since_epoch))

        if self.pcapng_port is not None and self.HCI_to_pcapng:
            self.pcapng_port.sendall(pcapng_h4_phdr(self.pcapng_interfaces[com_idx], True, payl, us_since_epoch))


    def read_timeout(self, limit = 0.5):
        if limit is not None and limit > 0:
            start = time.time()
            while time.time() < start + limit:
                self.read_all()

    def send_cmd_ints(self, com_idx, L, read_after = 2.0):
        payl = bytearray()
        for b in L:
            payl += s.pack('B', b)

        self.send_cmd(com_idx, payl)
        self.read_timeout(read_after)


    def send_parallel_cmd_seqs(self, lists, coms = (0,1), read_after = 0.05):
        for i in range(max([len(l) for l in lists])):
            for com_idx, cmds in zip(coms, lists):
                if i < len(cmds):
                    self.send_cmd(com_idx, cmds[i])

            self.read_timeout(read_after)

    def send_cmd_sequence(self, com_idx, cmds, read_after = 0.1):
        for cmd in cmds:
            self.send_cmd(com_idx, cmd)
            if read_after is not None:
                self.read_timeout(read_after)

    def read_version(self, com_idx):
        self.send_cmd(com_idx, b'\x01\x01\x10\x00')

    def read_all(self):
        for idx in range(len(self.coms)):
            ret = self.coms[idx].read(10000)
            if len(ret):
                self.handle_input(idx, ret)
    
    def handle_bytes(self, com_idx):
        return hp(s.pack('<H', self.conn_handles[com_idx]))

    def log(self, msg, sev = 0):
        if self.print_logs:
            print('\n{}'.format(msg))

        if self.pcapng_port:
            l = ceil((len(msg) + 1)/ 4) * 4 + 24
            extra =  l - (len(msg) + 1 + 24)
            ts_msb, ts_lsb = pcapng_ts()
            self.pcapng_port.sendall(s.pack('<IIIIIB', 0x80000109, l, ts_msb, ts_lsb, len(msg), sev) + bytes(msg, 'utf-8') + bytes(extra) + s.pack('<I', l))

    def on_connection_complete(self, com_idx, pkt):
        if not pkt.hdr_matches([4, 0x3e, 0x1f, 0x0a, 0x00]) and not pkt.hdr_matches([0x04, 0x3e, 0x22, 0x29, 0x00]):
            return

        self.conn_handles[com_idx] = pkt.bytes[5] | (pkt.bytes[6] << 8)

    def on_LTK_request(self, com_idx, pkt):
        if not pkt.hdr_matches([4, 0x3e, 0xd, 0x5, self.conn_handles[com_idx] & 0xFF, self.conn_handles[com_idx] >> 8]):
            return

        self.log('Sending LTK to board {}; responding to packet {}'.format(com_idx, pkt))
        self.encrypted = True

        self.send_cmd_sequence(com_idx, respond_to_connection_cmds(self.conn_handles[com_idx], self.LTK), None)

    def on_read_BD_ADDR_rsp(self, com_idx, pkt):
        if not pkt.hdr_matches([0x04, 0x0e, 0x0a, 0x01, 0x09, 0x10, 0x00]):
            return
        
        b = pkt.bytes[7:]
        addr = b[5] | (b[4] << 8) | (b[3] << 16) | (b[2] << 24) | (b[1] << 32) | (b[0] << 40)
        if com_idx == 0:
            self.cen_addr = addr
        else:
            self.adv_addr = addr

    def set_random_addresses(self):
        self.send_cmd_sequence(0, set_random_addr(self.cen_addr))
        self.send_cmd_sequence(1, set_random_addr(self.adv_addr))

    def add_to_filter_accept_list(self):
        self.send_cmd_sequence(0, add_to_filter_accept(self.adv_addr))

    def get_BD_ADDR(self, com_idx):
        self.input_cbs.append(self.on_read_BD_ADDR_rsp)
        self.send_cmd_sequence(com_idx, read_BD_ADDR_cmd(), 0.5)

    def on_encryption_v1(self, com_idx, pkt):
        if not pkt.hdr_matches([4, 8, 4, 0, self.conn_handles[com_idx] & 0xFF, self.conn_handles[com_idx] >> 8, 1]):
            return

        self.log('Board {}: Encryption enabled'.format(com_idx))
        self.encrypted_count = self.encrypted_count + 1        

    def connect(self, timeout = 10, debug = True):
        self.conn_handles = [None, None]
        self.send_init_cmds()
        if self.random_addresses:
            self.set_random_addresses()

        self.add_to_filter_accept_list()

        self.send_cmd_sequence(1, adv_cmds(self.random_addresses), 0.5)
        if debug:
           self.log('Board 1 Advertising\n')

        self.input_cbs.append(self.on_connection_complete)
        start = time.time()
        self.send_cmd_sequence(0, connect_cmd(self.random_addresses, self.adv_addr))
        while self.conn_handles[0] is None and time.time() < start + timeout:
            self.read_all()
            time.sleep(0.2)

        if self.conn_handles[0] is None:
            self.send_cmd_sequence(0, connect_cmd(self.random_addresses, self.adv_addr))
            while self.conn_handles[0] is None and time.time() < start + timeout:
                self.read_all()
                time.sleep(0.2)

        if debug:
            self.log('Connected\n')

    def disconnect(self, timeout = 2):
        start = time.time()
        self.send_cmd_sequence(0, disconnect_cmd(self.conn_handles[0]))
        while time.time() < start + timeout:
            self.read_all()
            time.sleep(0.2)

    def encrypt(self, timeout = 10):
        if self.conn_handles[0] is None:
            print("Cannot encrypt, no connection handle")
            return
        
        self.encrypted_count = 0
        self.input_cbs.append(self.on_LTK_request)
        self.input_cbs.append(self.on_encryption_v1)
        if self.pcapng_port is not None:
            self.pcapng_port.sendall(pcapng_key_pkt(self.LTK))

        self.send_cmd_sequence(0, encrypt_cmd(self.conn_handles[0], self.LTK))
        # self.send_cmd_sequence(0, encrypt_cmd(self.conn_handle, self.LTK))
        start = time.time()
        while self.encrypted_count < 2 and time.time() < start + timeout:
            self.read_all()

        time.sleep(0.3)


if __name__ == '__main__':
    pass
