import socket as skt
from math import ceil
import time
import struct

class blueSPY_socket:
    def __init__(self, port = 24352, IP="127.0.0.1"):
        self.port = port
        self.IP = IP
        self.socket = skt.socket(skt.AF_INET, skt.SOCK_STREAM)
        try:
            self.socket.connect((self.IP, self.port))
        except ConnectionRefusedError as e:
            print("Failed to connecct: {}/{}", self.IP, self.port)
    
    def send(self, data):
        self.socket.sendall(data)

def pcapng_ts(us_since_epoch = None):
    if us_since_epoch is None:
        us_since_epoch = int(round(time.time() * 1000000))

    ns_since_epoch = us_since_epoch * 1000

    return ns_since_epoch >> 32, 0xFFFFFFFF & ns_since_epoch

def LSB128(i):
    return (2**64 - 1) & i

def MSB128(i):
    return i >> 64

def pcapng_interface_description(link_type = 201, description = ''):
    opts_len = 0
    l = 20 + opts_len
    return struct.pack('<IIHHII', 0x1, l, link_type, 0xFFFF, 0xFFFFFFFF, l)

def pcapng_enhanced_pkt(id, payl, us_since_epoch = None):
    ts_msb, ts_lsb = pcapng_ts(us_since_epoch)

    l = 32 + 4 * ceil(len(payl) / 4)    
    extra =  l - (len(payl) +  32)
    return (struct.pack('<IIIIIII', 0x6, l, id, ts_msb, ts_lsb, len(payl), len(payl)) + 
                    bytes(payl) + bytes(extra) + struct.pack('<I', l))

def pcapng_h4_phdr(id, H2C, payl, us_since_epoch = None):
    if H2C:
        direction = 0
    else:
        direction = 1
    
    payl = struct.pack('>I', direction) + bytes(payl)
    return pcapng_enhanced_pkt(id, payl, us_since_epoch)

def log(msg, sev = 0):
    l = ceil((len(msg) + 1)/ 4) * 4 + 24
    extra =  l - (len(msg) + 1 + 24)
    ts_msb, ts_lsb = pcapng_ts()
    return (struct.pack('<IIIIIB', 0x80000109, l, ts_msb, ts_lsb, len(msg), sev) + bytes(msg, 'utf-8') + bytes(extra) + struct.pack('<I', l))

def pcapng_link_key(key, label: str = 'TCP Key', addr0 = 0, addr1 = 0):
    length = 40 + 4 * ceil((1 + len(label)) / 4)
    extra =  length - (len(label) + 1 + 40)
    return (struct.pack('<IIQQIHHIB', 0x80000200, length, LSB128(key), MSB128(key), 
                   0xFFFFFFFF & addr0, addr0 >> 32, 0xFFFF & addr1, addr1 >> 16, len(label)) + 
                bytes(label, 'utf-8') + bytes(extra) + struct.pack('<I', length))



if __name__ == '__main__':
    sock = blueSPY_socket()
    l_key = pcapng_link_key(0x5186976B321F2BD1AFE7334987E44B9A)
    sock.send(l_key) 