# Copyrite IBM 2022, 2025
# IBM Confidential
# Assisted by watsonx Code Assistant

# CTAP2HID reads bytes from endpoint 1,building the frames into a CTAPMsg. A CTAPMsg is used to initalize a 
# FIDO2Transaction thread which will reply on endpoint 2 when the transaction is complete.
#
# Authenticator uses some environmental properties to read PKI from the filesystem.
#       FIDO_AUTHENTICATOR_PKCS12 :: pkcs12 which contains the CA key/x509 and any generated x509 by the 
#                                    Fido2Authenticator.
#       FIDO_AUTHENTICATOR_PKCS12_SECRET :: password to decrypt the pkcs12 bag.
# Any credentials crated using the CA pkcs12 can be reconstructed during `navigator.credential.get` requests.
#       FIDO2_GENERATE_CA :: boolean to indicate if an error should be thrown if the PKCS12 file does not exist or
#                            cannot be read / converted to requried PKI. If this value is true the CA key and x509
#                            are generated on the first request to create a credential.
#

import base64, datetime, multiprocessing, os, sys, random, queue, threading, time, secrets, traceback
import cbor2 as cbor
from enum import Enum, IntEnum
from cryptography import x509
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import hashes
from cryptography.fernet import Fernet
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from usb_ip import BaseStructure, USBDevice, InterfaceDescriptor, DeviceConfigurations, EndPoint, USBContainer, \
                    USBIPCMDSubmit, bcolors, dump_bytes, colour_print
from key_pair import KeyPair, KeyUtils
from authenticator import Fido2Authenticator
from cert_utils import CertUtils


# Classes to create a USB device object, connect it to the system usb daemon and begin tx/rx
class CTAP2HIDClass(BaseStructure):
    _fields_ = [
        ('bLength', 'B', 9),
        ('bDescriptorType', 'B', 0x21),  # HID
        ('bcdHID', 'H'),
        ('bCountryCode', 'B'),
        ('bNumDescriptors', 'B'),
        ('bDescriptprType2', 'B'),
        ('wDescriptionLength', 'H'),
    ]


hid_class = CTAP2HIDClass(bcdHID=0x0100,  # HID version number
                     bCountryCode=0x0,
                     bNumDescriptors=0x1,
                     bDescriptprType2=0x22,  # Report
                     wDescriptionLength=0x3F00)  # Little endian


interface_d = InterfaceDescriptor(bAlternateSetting=0,
                                  bNumEndpoints=2,
                                  bInterfaceClass=3,  # class HID
                                  bInterfaceSubClass=0, # no interface subclass
                                  bInterfaceProtocol=0, # no interface protocol
                                  iInterface=0)

end_point_one = EndPoint(bEndpointAddress=0x04, # HOST OUT / USB IN
                     bmAttributes=0x3, # Interrupt transfer
                     wMaxPacketSize=(64&0x00FF)<<8 | (64&0xFF00),  # 64-byte packet max
                     bInterval=5)  # Poll every 5 millisecond

end_point_two = EndPoint(bEndpointAddress=0x8E, # HOST IN / USB OUT
                     bmAttributes=0x3, # Interrupt transfer
                     wMaxPacketSize=(64&0x00FF)<<8 | (64&0xFF00),  # 64-byte packet max, bussit
                     bInterval=5)  # Poll every 5 millisecond

interface_d.descriptions = [hid_class]  # Supports only one description
interface_d.endpoints = [end_point_two, end_point_one]  # Supports two endpoint
# wTotalLength = len( interface_d.pack() + hid_class.pack() + end_point_one.pack() + end_point_two.pack() + 18)
configuration = DeviceConfigurations(wTotalLength=0x2900,
                                     bNumInterfaces=0x1,
                                     bConfigurationValue=0x1,
                                     iConfiguration=0x0,  # No string
                                     bmAttributes=0x80,  # valid self powered
                                     bMaxPower=50)  # 50 mah current
configuration.interfaces = [interface_d]   # Supports only one interface


class Authenticator(object):
    '''
    Authenticator wraps the Fido2Authenticator with management for a persistent key/ca certificate stored in
    an encrypted file.

    The default location for this is in the ~/.fido2 directory for the user running this module, however this
    can be overwritten by environment properties.

    All keys generated by this authenticator are created so that the credential id is the encrypted key
    used to generate the attestation. This is done to allow this Authenticator to restore any key created with
    its ca key / AES key combination without the need to persist every individual private key used to 
    generate a new credential.
    '''

    _open_keys = {}
    _lock = multiprocessing.Lock()

    _pin_token_kp = None
    _pin_retry = 5

    def __new__(cls):
        cls._lock.acquire()
        if cls._pin_token_kp == None:
            cls._pin_token_kp = KeyPair.generate_ecdsa() #P-256 key, cose key type -25
        cls._lock.release()

    @classmethod
    def _validate_pin(cls, pinHash, cid):
        '''
        Look in the home direcotry of the user for a .fido2 directory. If it exists and contains
        enough data, try and open it.

        If we can open a given file and it contains the cert/key we expect; add it to cls._open_keys with the 
        pinAuthToken as the key.

        If we successfuly open and validata a file, then a generated pinAuthToken, which will
        function as an opaque lookup key, is returned.
        If no file could be opened then this method will return None
        '''
        # Try all .passkey files in the $HOME/.fido2 dir. If we parse out a cert/key then we have a 
        # valid pin and can cache the result and return an auth token.
        aesKey = algorithms.AES128(pinHash)
        baseDir = os.path.realpath(os.path.join(os.environ.get("HOME"), '.fido2'))
        for maybeKeyFile in os.listdir(baseDir):
            if not maybeKeyFile.endswith('.passkey'):
                colour_print(colour=bcolors.WARNING, component='Authenticator_validate_pin', 
                             msg='{} has invalid file type'.format(maybeKeyFile))
                continue
            encBagF = None
            try:
                encBagF = os.path.join(baseDir, maybeKeyFile)
                passkey = KeyUtils._load_passkey(pinHash, encBagF)
                colour_print(colour=bcolors.OKPINK, component='Authenticator_validate_pin', 
                             msg='Pin decrypted a .passkey file')
                ca = CertUtils.load_der_certificate(passkey.get('x5c'))
                kp = KeyUtils.load_der_key(passkey.get('key'))
                seed = passkey.get('seed')
                cls._pin_retry = 5 #Reset the counter
                pinAuthToken = secrets.token_bytes(32)
                cls._open_keys[cid] = {
                            'fKey': Fernet(seed),
                            'pem': ca,
                            'kp': kp,
                            'file': encBagF,
                            'ph': pinHash
                        }
                return pinAuthToken
            except Exception as e:
                colour_print(colour=bcolors.WARNING, component='Authenticator_validate_pin', 
                             msg='Something bad happened:\n{}'.format(e))
                continue
        colour_print(colour=bcolors.FAIL, component='Authenticator_validate_pin', 
                     msg='No valid pin found!')
        return None

    @classmethod
    def get_pin_cose_key(cls, pin_req, cid):
        if cls._pin_token_kp == None:
            cls.__new__(cls)
        return {1: KeyUtils.get_cose_key(cls._pin_token_kp.get_public(), hashes.SHA256(), ecdh=True)}

    @classmethod
    def get_pin_retries(cls, pin_req, cid):
        cls._pin_retry -= 1
        return {3: cls._pin_retry}

    @classmethod
    def decapsulate(cls, ecCoseKey):
        cose_type_to_curve_map = { #These are kind of made up, as per 
                    #https://fidoalliance.org/specs/fido-v2.1-ps-20210615/fido-client-to-authenticator-protocol-v2.1-ps-errata-20220621.html#pinProto1
                    -25: ec.SECP256R1,
                    -26: ec.SECP521R1
                }
        ec_pub_numbs = ec.EllipticCurvePublicNumbers(KeyUtils._bytes_to_long(ecCoseKey[-2]),
                            KeyUtils._bytes_to_long(ecCoseKey[-3]), 
                            cose_type_to_curve_map[ecCoseKey[3]]())
        pubkey = ec_pub_numbs.public_key()
        shared_point = cls._pin_token_kp.get_private().exchange(ec.ECDH(), pubkey)
        hasher = hashes.Hash(hashes.SHA256());
        hasher.update(shared_point)
        return hasher.finalize()

    @classmethod
    def get_pin_token(cls, pin_req, cid):
        #https://fidoalliance.org/specs/fido-v2.1-ps-20210615/fido-client-to-authenticator-protocol-v2.1-ps-errata-20220621.html#getPinToken
        platform_cose_key = pin_req[3]
        pin_hash_enc = pin_req[6]
        colour_print(colour=bcolors.OKPINK, component='Authenticator.get_pin_token', 
                     msg='plat cose key: {}; pinHashEnc: {}'.format(platform_cose_key, pin_hash_enc))
        sharedSecret = cls.decapsulate(platform_cose_key)
        colour_print(colour=bcolors.OKPINK, component='Authenticator.get_pin_token', 
                     msg='shared secret: {};'.format(sharedSecret))
        decryptor = Cipher(algorithms.AES256(sharedSecret), modes.CBC(bytes([0] * 16))).decryptor()
        pin = decryptor.update(pin_hash_enc) + decryptor.finalize()
        pinAuthToken = cls._validate_pin(pin, cid)
        if pinAuthToken != None:
            return {2: pinAuthToken}
        return None


    @classmethod
    def attestation_out(cls, clientDataHash, rp, user, pkCredsParams, excludeList, exts, cid):
        colour_print(colour=bcolors.OKPINK, component='Authenticator.attestation_out', 
                     msg='open keys: {}'.format(cls._open_keys))
        if not cid in cls._open_keys.keys():
            return None, None
        ca = cls._open_keys.get(cid)
        kp = KeyPair.generate_ecdsa()
        authenticator = Fido2Authenticator(keyPair=kp, caKeyPair=ca.get('kp'), 
                                            caCert=ca.get('pem'), fKey=ca.get('fKey'))
        credId = authenticator.get_credential_id(kp)
        pk = {'rp': rp} # authenticator.py expects this strucutre
        authData = authenticator.build_authenticator_data(pk, 'packed', kp, True, up=True, be=False, bs=False)
        colour_print(colour=bcolors.OKPINK, component='Authenticator.attestation_out', 
                     msg='credId: {}; toSign: {}'.format(credId, base64.b64encode(bytes([*authData, *clientDataHash ])).decode()))
        attStmt = authenticator.build_packed_attestation_statement('packed', 
                                                    clientDataHash, authData, None, kp)
        colour_print(colour=bcolors.OKPINK, component='Authenticator.attestation_out', 
                     msg='attStmt: {}'.format(attStmt))
        #Since we set UV we can make these resident keys
        KeyUtils.update_passkey({rp['id']: credId, 'user': user['id']}, ca['ph'], ca['file'])
        return authData, attStmt


    @classmethod
    def assertion_out(cls, rpId, clientDataHash, allowedList, exts, cid):
        if not cid in cls._open_keys.keys():
            return None, None, None, None
        ca = cls._open_keys.get(cid)
        resCreds = KeyUtils._load_passkey(ca['ph'], ca['file']).get('res_creds')
        if resCreds != None:
            colour_print(colour=bcolors.OKPINK, component='FIDO2Authenticator.assertion_outputs',
                         msg='passkey has resident credentials, adding them to allowed list')
            for rc in resCreds:
                if rc.get(rpId) != None:
                    allowedList += [{'id': base64.urlsafe_b64decode(rc[rpId]), 
                                     'user': rc.get('user')}]
        for cred in allowedList:
            try:
                ca_pem = ca.get('pem')
                ca_kp = ca.get('kp')
                fKey = ca.get('fKey')
                #Create the authenticator, creating the key pair will fail if we don'w own the credId
                b64CredId = base64.urlsafe_b64encode(cred.get('id'))
                kp = Fido2Authenticator._get_key_pair_from_credential_id(b64CredId, fKey)
                colour_print(colour=bcolors.OKPINK, component='FIDO2Authenticator.assertion_outputs',
                             msg='We have a usable key, sign the challenge')
                _authenticator = Fido2Authenticator(keyPair=kp, credId=b64CredId, aaguid=[b'\x00'*16], 
                                                    caKeyPair=ca_kp, caCert=ca_pem, fKey=fKey)
                credential = {
                        "id": _authenticator.cib, #Should be set by __init__()
                        "type" : "public-key"
                    }
                #Generate the assertion response data
                authData = _authenticator.build_authenticator_data({'rpId': rpId}, 'packed', kp, True, 
                                                                        up=True, be=False, bs=False)
                #Sign it, authenticator method b64 encodeds the ba so we decode it right away
                sig = _authenticator.assertion_signature(authData, clientDataHash, kp)
                userHandle = cred.get("user")
                return credential, authData, sig, userHandle
            except Exception as e:
                colour_print(colour=bcolors.FAIL, component='FIDO2Authenticator.assertion_outputs',
                             msg='Could not retrieve key pair from credential id {}'.format(cred))
                continue
        return None, None, None, None

class CBORCommand(object):

    class CommandByte(Enum):
        MAKE_CREDENTIAL = 0x1
        GET_NEXT_ASSERTION = 0x2
        GET_INFO = 0x4
        CLIENT_PIN = 0x6
        RESET = 0x7
        CREDENTIAL_MANAGEMENT = 0x9
        AUTHENTICATOR_SELECTION = 0xB
        AUTHENTICATOR_CONFIG = 0xD

        def __repr__(self):
            #return f'{self.__class__.__name__}.{self.name}'
            return self.value

    class CBORStatusCode(IntEnum):
        CTAP2_OK = 0x0
        CTAP1_ERR_INVALID_COMMAND = 0x01
        CTAP2_ERR_INVALID_CBOR = 0x12
        CTAP2_ERR_PIN_INVALID = 0x31
        CTAP2_ERR_PIN_AUTH_INVALID = 0x33
        CTAP2_ERR_PUAT_REQUIRED = 0x36


    cid = 0xFFFFFFFF
    request = []
    response = []
    response_segment = 0
    response_ready = False
    length = 0
    request_segment = 0
    cmd = None
    ctaphid_cmd = 0
    bcnt = 0

    def __init__(self, cid, ba, skip_init=False):
        self.cid = cid
        if ba == None and skip_init == True:
            return #Create an empty command as we will directly set the response buffer later with the assigned CID.
        if len(ba) <= 1:
            colour_print(colour=bcolors.OKYELLOW, component='CBORCommand.__init__', 
                    msg="Byte Array must be at least one byte long")
        # Length of the incoming CBOR message (total).
        self.length = int.from_bytes(ba[0:2], 'big')
        # Request buffer. This stores the incoming CBOR message and grows until all segments have been recieved
        self.request = ba[3:]
        # Track then number of response segments transmitted, the number transmitted in the continue sequence packet
        # should always be one less than this number
        self.response_segment = 0
        # Track the number of request segmetns recieved
        self.request_segment = 0
        # Response buffer. This stores the outgoing response to the recieved CBOR message and shrinks until the entire
        # response has been transmitted
        self.response = []
        # Command recieved in CTAPHID frame, this is likely 0x90 (CBOR_MSG) but might be different
        self.ctaphid_cmd = 0
        # Authenticator API command byte recieved in initial packet
        self.cmd = self.CommandByte(int.from_bytes(ba[2:3], 'big'))
        #Length of the payload bytes
        self.bcnt = 0
        # Signal that the response buffer is ready to be sent back to the client
        self.response_ready = False
        colour_print(colour=bcolors.OKPURPLE, component='CBORCommand.__init__', 
                msg="command {}; length {}; self.request[{}]".format(self.cmd, self.length, len(self.request)))
        if self.length >= len(self.request):
            colour_print(colour=bcolors.OKPURPLE, component='CBORCommand.__init__', 
                    msg="request is segmented, wait for the whole message")
        else: #We have the whole message
            self.response = self.unpack()
            dump_bytes(self.response, colour=bcolors.OKPINK, 
                       component='CBORCommand.__init__', msg='CTAP response')

    def append_segment(self, seg_buf):
        self.request += seg_buf
        if len(self.request) >= self.length:
            self.response = self.unpack()
            dump_bytes(self.response, colour=bcolors.OKPINK, 
                       component='CBORCommand.append_segment', msg='CTAP response')
        else:
            self.request_segment += 1


    #Return CBOR response if entire command has been received or None if still 
    #waiting for segments
    def unpack(self):
        return {
            self.CommandByte.MAKE_CREDENTIAL: self._make_cred,
            self.CommandByte.GET_NEXT_ASSERTION: self._get_assertion,
            self.CommandByte.GET_INFO: self._get_info,
            self.CommandByte.CLIENT_PIN: self._client_pin
            }.get(self.cmd)(bytes(self.request))

    def get_rsp_seg(self, num_bytes):
        if not isinstance(num_bytes, int):
            raise RuntimeError("panic!")
        self.response_segment += 1
        #sequence is offset by two to account for init pkt and zero index start for continue sequence
        seg_num = max(self.response_segment - 2, 0)
        colour_print(colour=bcolors.WARNING, component='CBORCommand.get_rsp_seg', msg='self.response_segment = {}'.format(self.response_segment))
        colour_print(colour=bcolors.WARNING, component='CBORCommand.get_rsp_seg', msg='self.response_segment - 2 = {}'.format(self.response_segment - 2))
        colour_print(colour=bcolors.WARNING, component='CBORCommand.get_rsp_seg', msg='segment number = {}'.format(seg_num))
        seg = self.response
        if num_bytes >= len(self.response):
            self.response = []
        else:
            seg = self.response[:num_bytes]
            self.response = self.response[num_bytes:]
        return seg, seg_num

    def _client_pin(self, ba):
        # https://fidoalliance.org/specs/fido-v2.2-rd-20230321/fido-client-to-authenticator-protocol-v2.2-rd-20230321.html#authenticatorClientPIN
        pin_sub_cmds = { 
                      1: Authenticator.get_pin_retries,
                      2: Authenticator.get_pin_cose_key,
            #SET_PIN = 0x3
            #CHANGE_PIN = 0x4
                      5: Authenticator.get_pin_token
                }
        req_data = cbor.loads(ba)
        colour_print(colour=bcolors.OKGREEN, component='CBORCommand._client_pin',
                     msg='Packet request: {}'.format(req_data))
        sub_cmd = req_data[2]
        colour_print(colour=bcolors.OKGREEN, component='CBORCommand._client_pin',
                     msg='pin sub_cmd: {}'.format(sub_cmd))
        rsp = pin_sub_cmds[sub_cmd](req_data, self.cid)
        result = (self.CBORStatusCode.CTAP2_ERR_PIN_INVALID).to_bytes()
        if rsp != None:
            result = (self.CBORStatusCode.CTAP2_OK).to_bytes() + cbor.dumps(rsp)
        self.bcnt = len(result)
        self.response_ready = True
        return result

    # authenticatorGetInfo takes no inputs so return immediately
    def _get_info(self, ba):
        result = {
            0x01: ["U2F_V2", "FIDO_2_0"],
            0x02: ['hmac-secret'],
            #0x03: b"\x01\x02\x03\x04" * 4,
            0x03: b"\x00" * 16,
            0x04: {'rk': True, 'up': True, 'plat': False, 'clientPin': True},
            0x05: 1200,
            0x06: [1]
        }
        result = (self.CBORStatusCode.CTAP2_OK).to_bytes() + cbor.dumps(result)
        self.bcnt = len(result)
        self.response_ready = True
        return result

    def _make_cred(self, ba):
        # https://fidoalliance.org/specs/fido-v2.2-rd-20230321/fido-client-to-authenticator-protocol-v2.2-rd-20230321.html#authenticatorMakeCredential
        req = cbor.loads(ba)
        colour_print(colour=bcolors.FAIL, component='CBORCommand._make_cred',
                     msg='CBOR request {}'.format(req))
        for prop in [(0x01, 'clientDataHash'), (0x02, 'rp'), (0x03, 'user'), (0x04, 'pubkeyCredParams')]:
            if not prop[0] in req.keys():
                colour_print(colour=bcolors.FAIL, component='CBORCommand._make_cred',
                             msg='{} missing from request:\n{}'.format(prop[1], cbor.dumps(req)))
                raise Exception("Missing required property %s" % prop[1])
        authData, attStmt = Authenticator.attestation_out(req.get(0x01), req.get(0x02), req.get(0x03),
                                            req.get(0x04), req.get(0x05), req.get(0x06), self.cid)
        result = (self.CBORStatusCode.CTAP2_ERR_PUAT_REQUIRED).to_bytes()
        if authData and attStmt:
            rsp = {
                0x01: 'packed', #fmt
                0x02: authData,
                0x03: attStmt
            }
            result = (self.CBORStatusCode.CTAP2_OK).to_bytes() + cbor.dumps(rsp)
        self.bcnt = len(result)
        self.response_ready = True
        return result


    def _get_assertion(self, ba):
        # https://fidoalliance.org/specs/fido-v2.2-rd-20230321/fido-client-to-authenticator-protocol-v2.2-rd-20230321.html#authenticatorGetAssertion
        req = cbor.loads(ba)
        colour_print(colour=bcolors.FAIL, component='CBORCommand._get_assertion',
                     msg='CBOR request {}'.format(req))
        for prop in [(0x01, 'rpId'), (0x02, 'clientDataHash')]:
            if not prop[0] in req:
                colour_print(colour=bcolors.FAIL, component='CBORCommand._make_cred',
                             msg='{} missing from request:\n{}'.format(prop[1], cbor.dumps(req)))
                raise Exception("Missing required property %s" % prop[1])
        credential, authData, signature, userHandle = Authenticator.assertion_out(req.get(0x01), 
                                                req.get(0x02), req.get(0x03, []), req.get(0x04, {}), self.cid)
        result = (self.CBORStatusCode.CTAP2_ERR_PUAT_REQUIRED).to_bytes()
        if credential and authData and signature:
            rsp = {
                    0x01: credential,
                    0x02: authData,
                    0x03: signature
            }
            if userHandle:
                rsp[0x04] = {'id': userHandle}
            result = (self.CBORStatusCode.CTAP2_OK).to_bytes() + cbor.dumps(rsp)
        self.bcnt = len(result)
        self.response_ready = True
        return result

# Send this for every new message frame. Response contains
class CTAPHIDInitPkt(BaseStructure):
    _fields_ = [
        ('cid', 'I'),
        ('cmd', 'B'),
        ('bcnt', 'H'),
    ]

    def __init__(self, **kwargs):
        if 'data' in kwargs:
            index = None
            for i, field in enumerate(self._fields_):
                if field[0] == 'data':
                    index = i
                    break
            if index == None:
                colour_print(colour=bcolors.OKGREEN, component='CTAPHIDInitPkt.__init__', msg='setting data field')
                self._fields_ += [('data', '%ds' % len(kwargs['data']))]
            else:
                colour_print(colour=bcolors.OKGREEN, component='CTAPHIDInitPkt.__init__', 
                             msg='data already exists as a field, updating it')
                self._fields_[index] = ('data', '%ds' % len(kwargs['data']))
            print(self._fields_)
        super().__init__(**kwargs)

class CTAPHIDSeqPkt(BaseStructure):

    _fields_ = [
        ('cid', 'I'),
        ('seq', 'B'),
    ]

    def __init__(self, **kwargs):
        print(kwargs)
        if 'data' in kwargs:
            index = None
            for i, field in enumerate(self._fields_):
                if field[0] == 'data':
                    index = i
                    break
            if index == None:
                print("setting data field")
                colour_print(colour=bcolors.OKPINK, component='CTAPHIDSeqPkt.__init__', msg='setting data field')
                self._fields_ += [('data', '%ds' % len(kwargs['data']))]
            else:
                colour_print(colour=bcolors.OKPINK, component='CTAPHIDSeqPkt.__init__', 
                             msg='data already exists as a field, updating it')
                self._fields_[index] = ('data', '%ds' % len(kwargs['data']))
            print(self._fields_)
        super(CTAPHIDSeqPkt, self).__init__(**kwargs)



class CTAP2HIDevice(USBDevice):
    speed = 2
    vendorID = 0x3713
    productID = 0x3713
    bcdDevice = 0x0
    bcdUSB = 0x0
    bNumConfigurations = 0x1
    bNumInterfaces = 0x1
    bConfigurationValue = 0x1
    bDeviceClass = 0x0
    bDeviceSubClass = 0x0
    bDeviceProtocol = 0x0
    bNumConfigurations = 1
    bConfigurationValue = 1
    bNumInterfaces = 1
    configurations = [configuration]  # Supports only one configuration

    #This will contain the current set of channel id's and associated state
    cids = {}
    #This is a list of pending message frames with a worker thread which will respond within 50ms if 
    # we don't generate a response in time
    pending = []

    def __init__(self):
        USBDevice.__init__(self)
        self.start_time = datetime.datetime.now()

    def unlink(self, usb_req):
        index = None
        for i, worker in enumerate(self.pending):
            if worker.req.seqnum == usb_req.unlink_seqnum:
                index = i
        if index != None:
            del self.pending[index]
            colour_print(component='CTAP2HIDevice.unlink', msg='unlinked request {}'.format(usb_req.unlink_seqnum))
            return True
        return False

    def generate_fido2_report(self):
        arr = [0x06, 0xd0, 0xf1,        # USAGE_PAGE (FIDO Alliance)
                0x09, 0x01,             # USAGE (CTAPHID)
                0xa1, 0x01,             # HID_Collection(HID_Application)
                0x09, 0x20,             # USAGE (Input Report Data)
                0x15, 0x00,             # LOGICAL_MINIMUM (0)
                0x26, 0xff, 0x00,       # LOGICAL_MAXIMUM (255)
                0x75, 0x08,             # REPORT_SIZE (8)
                0x95, 0x40,             # REPORT_COUNT (64)
                0x81, 0x02,             # INPUT (Data,Var,Abs)
                0x09, 0x21,             # USAGE(Output Report Data)
                0x15, 0x00,             # LOGICAL_MINIMUM (0)
                0x26, 0xff, 0x00,       # LOGICAL_MAXIMUM (255)
                0x75, 0x08,             # REPORT_SIZE (8)
                0x95, 0x40,             # REPORT_COUNT (64)
                0x91, 0x02,             # OUTPUT (Data,Var,Abs)
            0xc0]                   # End Collection
        return bytes(arr) 

    def _bytes_to_str(self, b):
        return ''.join("%02X" % x for x in b)


    class KeepAliveWorker(object):

        stop = False
        my_thread = None
        req = None
        hid = None
        start = None

        def __init__(self, req, hid):
            self.req = req
            self.hid = hid

        def reply_with_keepalive(self):
            start = round(time.time() * 1000)
            while self.stop == False:
                now = round(time.time() * 1000)
                wait_time = now - start
                if(now >= start + 450) and self.stop == False: #If we get to 45ms, send heartbeat
                    colour_print(colour=bcolors.FAIL, component='KeepAliveWorker.reply_with_keepalive',
                                 msg='Thread reached timeout of {} ms before response buffer was recieved . . . sending heartbeat response'.format(wait_time))
                    #We need to reply with keep alive after 50ms this gives us some tolerance
                    #Our status is always still processing
                    rsp = b'\x3B\x01\x01'
                    self.hid.send_usb_req(rsp, len(rsp), start_frame=0xFFFFFFFF, packets=0, ep=0, direction=0, 
                                       seqnum=self.req.seqnum)
                    self.stop = True
                    for idx, txn in enumerate(self.hid.pending): #If we sent a keep alive remove the transaction from the
                        # list of pending transactions
                        if txn.req is not None and txn.req.seqnum is not None and txn.req.seqnum == self.req.seqnum:
                            del self.hid.pending[idx]
                    return
                time.sleep(5/1000) # sleep for 5 ms


    def send_response_segment(self, cid, cbor_cmd):
        if len(self.pending) == 0:
            colour_print(colour=bcolors.FAIL, component='send_response_segment', msg='No pending transactions to use :(')
            return
        #Take control of the next pending request, if this fails then we will just have to wait for the
        # next frame to appear.
        self.pending[0].stop = True
        if self.pending[0].my_thread is not None and self.pending[0].my_thread.is_alive() == True:
            self.pending[0].my_thread.join(1/1000)
        if self.pending[0].my_thread.is_alive() == True:
            colour_print(colour=bcolors.FAIL, component='send_response_segment', msg='Could not kill keepalive thread')
            return
        rsp_data = None
        if cbor_cmd.response_segment == 0: #We send the init pkt
            data, _ = cbor_cmd.get_rsp_seg(57)
            rsp_data = CTAPHIDInitPkt(cid=int.from_bytes(cid, 'big'), 
                                    cmd=cbor_cmd.ctaphid_cmd,
                                    bcnt=cbor_cmd.bcnt,
                                    data=data).pack()
        else: #We send the continue sequence pkt
            data, seq_num = cbor_cmd.get_rsp_seg(59)
            colour_print(colour=bcolors.WARNING, component='send_response_segment', msg='Sequence number {}'.format(seq_num))
            rsp_data = CTAPHIDSeqPkt(cid=int.from_bytes(cid, 'big'),
                                     seq=seq_num,
                                     data=data).pack()
        colour_print(colour=bcolors.WARNING, component='send_response_segment', msg='pad with {} 0 bytes'.format(64 - len(rsp_data)))
        rsp_data += b'\00' * (64 - len(rsp_data)) # pad the 64 byte frame with 0x00 if required
        '''
        rsp_data = CTAPHIDInitPkt(cid=int.from_bytes(cid), cmd=int.from_bytes(cmd), bcnt=17, data=data).pack()
        dump_bytes(rsp, colour=bcolors.OKGREEN, component='USBDevice.ctaphid_init', msg='Packed response')
        rsp_data = cid
        rsp_data += int.to_bytes(cbor_cmd.response_segment, 1)
        if len(cbor_cmd.response) > 58:
            rsp_data += cbor_cmd.response[:58]
            cbor_cmd.response = cbor_cmd.response[58:]
        else:
            rsp_data += cbor_cmd.response
            cbor_cmd.response = []
        # Pad out to 64 bytes
        rsp_data += b'\x00' * (64 - len(rsp_data))
        '''
        dump_bytes(rsp_data, colour=bcolors.OKGREEN, component='USBDevice.send_response_segment', 
                   msg='Packed response: ')
        self.send_usb_req(rsp_data, len(rsp_data), start_frame=0xFFFFFFFF, packets=0, 
                          ep=0, direction=0, seqnum=self.pending[0].req.seqnum)
        del self.pending[0]
        #If we have not response buffer left, remove the transaction from our context
        # don't do this for the broadcast channel.
        if len(cbor_cmd.response) == 0 and cid != b'\xff\xff\xff\xff':
            del self.cids[cid]['cborCmd']
        return


    def ctaphid_ping(self, usb_req):
        return

    def ctaphid_msg(self, usb_req):

        class U2FCommand(Enum):
            U2F_VERSION = 0x0
            U2F_REGISTER = 0x1
            U2F_AUTHENTICATE = 0x2
            U2F_VER = 0x03

        #Only supporting extended length encoding, section 3.1.3
        #https://fidoalliance.org/specs/fido-u2f-v1.2-ps-20170411/fido-u2f-raw-message-formats-v1.2-ps-20170411.html
        cid = usb_req.data_frame[0:4]
        if not cid in self.cids:
            colour_print(colour=bcolors.FAIL, component='USBDevice.ctaphid_msg', 
                         msg='Unknown CID {}'.format(cid))
        cmd = usb_req.data_frame[4:5]
        bcnt = usb_req.data_frame[5:7]
        apdu = usb_req.data_frame[7:]
        colour_print(colour=bcolors.OKGREEN, component='USBDevice.ctaphid_msg', 
                    msg='cmd = {}; bcnt = {}; apdu = {}'.format(
                        self._bytes_to_str(cmd), self._bytes_to_str(bcnt), apdu))
        u2f_cla = apdu[:1]
        u2f_ins = apdu[1:2]
        u2f_p1 = apdu[2:3]
        u2f_p2 = apdu[3:4]
        u2f_lc = apdu[4:7]
        u2f_data = apdu[7:]
        colour_print(colour=bcolors.OKGREEN, component='USBDevice.ctaphid_msg', 
                    msg='U2F Raw Message CLA = {}; INS = {}; P1 = {}; P2 = {}; Lc = {}'.format(
                        u2f_cla, u2f_ins, u2f_p1, u2f_p2, u2f_lc))
        if u2f_cla == b'\x00' and u2f_ins == b'\x03':
            #U2F_VERSION request, send the expected response
            cborCmd = CBORCommand(cid, None, skip_init=True)
            cborCmd.ctaphid_cmd = int.from_bytes(cmd, 'big')
            cborCmd.bcnt = 6
            cborCmd.response = b'U2F_V2'
            self.cids[cid]['cborCmd'] = cborCmd
            self.send_response_segment(cid, self.cids[cid]['cborCmd'])

    def ctaphid_init(self, usb_req):
        cid = usb_req.data_frame[0:4]
        cmd = usb_req.data_frame[4:5]
        bcnt = usb_req.data_frame[5:7]
        nonce = usb_req.data_frame[7:15]
        assignedCID = bytes([0, random.randint(0, 255), 0, random.randint(0, 255)])
        colour_print(colour=bcolors.OKGREEN, component='USBDevice.ctaphid_init', 
                    msg='Assigning a new CID to {}'.format(self._bytes_to_str(assignedCID)))
        data = nonce + assignedCID
        # protocol == 2; major version == 5; minor version = 2; build version = 5
        # data += int.to_bytes(2) + int.to_bytes(5) + int.to_bytes(1) + int.to_bytes(2) + int.to_bytes(5)
        for i in [2, 5, 1, 2, 0x04 | 0x08]:
            data += int.to_bytes(i)
        dump_bytes(data, colour=bcolors.OKGREEN, component='USBDevice.ctaphid_init', msg='Response data')
        data += b'\00' * (57 - len(data)) # 64 - 4 (CID) - 1 (cmd) - 2 (bcnt) - len of response
        #rsp = CTAPHIDInitPkt(cid=int.from_bytes(cid), cmd=int.from_bytes(cmd), bcnt=17, data=data).pack()
        #dump_bytes(rsp, colour=bcolors.OKGREEN, component='USBDevice.ctaphid_init', msg='Packed response')
        self.cids[assignedCID] = {'cborCmd': CBORCommand(cid, None, skip_init=True) }
        self.cids[assignedCID]['cborCmd'].response = data
        self.cids[assignedCID]['cborCmd'].ctaphid_cmd = int.from_bytes(cmd, 'big')
        self.cids[assignedCID]['cborCmd'].bcnt = 17
        #return self.send_usb_req(usb_req, rsp, 64, ep=0x81, start_frame=int.from_bytes(cid))
        #self.send_usb_req(usb_req, b'', 0, ep=usb_req.ep)
        #self.send_usb_req(rsp, len(rsp), start_frame=0xFFFFFFFF, packets=0, 
        #                  ep=0, direction=0, seqnum=self.pending[0].seqnum)
        #del self.pending[0]
        self.send_response_segment(cid, self.cids[assignedCID]['cborCmd'])

    def ctaphid_cbor(self, usb_req):
        cid = usb_req.data_frame[0:4]
        colour_print(colour=bcolors.OKGREEN, component='USBDevice.ctaphid_cbor', 
                    msg='CBOR message recieved on channel {}'.format(self._bytes_to_str(cid)))
        cmd = usb_req.data_frame[4:5]
        bcnt = usb_req.data_frame[5:7]
        ctap_cmd = usb_req.data_frame[7:8]
        print(int.from_bytes(bcnt, 'big') - 1)
        cbor_data = usb_req.data_frame[8: 7 + int.from_bytes(bcnt, 'big')]
        colour_print(colour=bcolors.OKGREEN, component='USBDevice.ctaphid_cbor', 
                     msg='CBOR msg frame cmd: {}; bcnt: {}'.format(self._bytes_to_str(ctap_cmd),
                                                                   self._bytes_to_str(bcnt)))
        dump_bytes(cbor_data, colour=bcolors.OKGREEN, component='USBDevice.ctaphid_cbor', 
                    msg='CBOR encoded bytes: ')
        cbor_cmd = CBORCommand(cid, usb_req.data_frame[5:])
        cbor_cmd.ctaphid_cmd = int.from_bytes(cmd, 'big')
        self.cids[cid]['cborCmd'] = cbor_cmd
        if cbor_cmd.response_ready == True: #We can respond immediatly
            dump_bytes(self.cids[cid]['cborCmd'].response, colour=bcolors.OKGREEN, 
                       component='USBDevice.ctaphid_cbor', msg='CBOR response: ')
            return self.send_response_segment(cid, self.cids[cid]['cborCmd'])
        else:
            colour_print(colour=bcolors.OKYELLOW, component='USBDevice.ctaphid_cbor', 
                         msg="Waiting for rest of command to arrive . . .")
            return

    def ctaphid_cancel(self, usb_req):
        return

    def ctaphid_keepalive(self, usb_req):
        return

    def ctaphid_error(self, usb_req):
        return

    def ctaphid_unknown(self, usb_req):
        colour_print(colour=bcolors.FAIL, component='USBDevice.ctaphid_unknown', msg='Unkown request recieved')
        self.send_usb_req(b'', 0, ep=usb_req.ep, seqnum=usb_req.seqnum)

    def _handle_outgoing(self, usb_req):
        colour_print(component='USBDevice._handle_outgoing', msg='Adding request to pending')
        # Create a new thread so we will reply at 50ms if required
        worker = self.KeepAliveWorker(usb_req, self)
        keep_alive_thread = threading.Thread(target=worker.reply_with_keepalive)
        worker.my_thread = keep_alive_thread
        #keep_alive_thread.start()
        self.pending.append( worker )
        for cid, context in self.cids.items():
            if context.get('cborCmd') != None and context['cborCmd'].response_ready == True:
                colour_print(component='USBDevice._handle_outgoing', 
                            msg='Using pending request to send response segment')
                self.send_response_segment(cid, context['cborCmd'])
                return

    def _handle_incoming_cmd(self, cmd, usb_req):
        ctapCmd = int.from_bytes(cmd, 'big') & 0x7F
        colour_print(colour=bcolors.OKGREEN, component='USBDevice._handle_incoming_cmd', 
                    msg='recieved command {}'.format(ctapCmd))
        return {
            1: self.ctaphid_ping,
            3: self.ctaphid_msg,
            6: self.ctaphid_init,
            16: self.ctaphid_cbor,
            17: self.ctaphid_cancel,
            59: self.ctaphid_keepalive,
            63: self.ctaphid_error,
        }.get(ctapCmd, self.ctaphid_unknown)(usb_req)

    def _handle_incoming_sequence(self, cid, usb_req):
        seqNum = int.from_bytes(usb_req.data_frame[4:5], 'big')
        context = self.cids.get(cid)
        if context != None:
            transaction = context.get("cborCmd")
            if transaction != None and seqNum == transaction.request_segment:
                transaction.append_segment(usb_req.data_frame[5:])
                if transaction.response_ready == True:
                    self.send_response_segment(cid, transaction)
                else:
                    colour_print(colour=bcolors.OKPURPLE, component='USBDevice._handle_incoming_sequence', 
                                 msg='Sequence number [{}] not the last expected sequence'.format(seqNum))
            else:
                colour_print(colour=bcolors.FAIL, component='USBDevice._handle_incoming_sequence', 
                             msg='Sequence number [{}] not the next expected sequence [{}]'.format(
                                 seqNum, transaction.request_segment))
        else:
            colour_print(colour=bcolors.FAIL, component='USBDevice._handle_incoming_sequence', 
                         msg='CID not found in device context, don\'t know what to do')

    def _handle_incoming(self, usb_req):
        if len(self.pending) == 0:
            colour_print(colour=bcolors.FAIL, component="USBDevice._handle_incoming", 
                         msg="No pending request to respond with :(")
        else: #We reply to the HOST Out endpoint with an empty frame then process the command
            self.send_usb_req('', 64, ep=0, start_frame=0xFFFFFFFF, seqnum=usb_req.seqnum)
        # We send the response to the command to a request in self.pending once we are ready
        cid = usb_req.data_frame[0:4]
        cmd = usb_req.data_frame[4:5]
        colour_print(colour=bcolors.OKGREEN, component='USBDevice._handle_incoming', 
                    msg='CID: {}; command: {}'.format(self._bytes_to_str(cid), self._bytes_to_str(cmd)))
        if(int.from_bytes(cmd, 'big') & 0x80) > 0:
            colour_print(colour=bcolors.FAIL, component='USBDevice._handle_incoming', 
                        msg='bit 8 set we got a command msg')
            return self._handle_incoming_cmd(cmd, usb_req)
        else:
            colour_print(colour=bcolors.OKPURPLE, component='USBDevice._handle_incoming',
                         msg='Recieved a sequence segment, appending it to the current msg context')
            return self._handle_incoming_sequence(cid, usb_req)

    def handle_data(self, usb_req):
        '''
        if usb_req.start_frame == 0xFFFFFFFF:
            colour_print(component='USBDevice.handle_data', msg='Don\'t know what to do with this request')
            self.send_usb_req(b'', 0, seqnum=usb_req.seqnum)
        '''
        if usb_req.ep == 0xE: #HOST In endpoint, add it to the queue so we can use it when we have data 
            return self._handle_outgoing(usb_req) # or use it to send the current in progress transaction
        else: #We have data, work out what to do with it and if we have a pending request we can respond with
            return self._handle_incoming(usb_req)


    def handle_unknown_control(self, control_req, usb_req):
        handled = False
        if control_req.bmRequestType == 0x81: #Interface request
            if control_req.bRequest == 0x6:  # Get Descriptor
                if control_req.wValue == 0x22:  # send initial report
                    print('[' + bcolors.OKGREEN + 'USBDevice' + bcolors.ENDC + '] Send initial report ')
                    ret=self.generate_fido2_report()
                    self.send_usb_req(ret, len(ret), seqnum=usb_req.seqnum)
                    handled = True
        elif control_req.bmRequestType == 0x21:  # Host Request
            if control_req.bRequest == 0x0a:  # set idle
                print('[' + bcolors.OKGREEN + 'USBDevice' + bcolors.ENDC + '] HID Idle ')
                self.send_usb_req(b'', 0, seqnum=usb_req.seqnum)
                handled = True
        else:
            print('[' + bcolors.FAIL + 'USBDevice' + bcolors.ENDC + '] Unknown control [{}] '.format(
                ', '.join( hex(x) for x in list(usb_req.setup.to_bytes(8, 'big')) )))
            print('[' + bcolors.FAIL + 'USBDevice' + bcolors.ENDC + '] Unknown flags [{}] '.format(
                ', '.join( hex(x) for x in list(usb_req.flags.to_bytes(8, 'big')) )))
            print('[' + bcolors.FAIL + 'USBDevice' + bcolors.ENDC + '] Unknown data [{}] '.format(
                ', '.join( hex(x) for x in list(usb_req.data_frame) )))
            #raise RuntimeError("unknown control")
            self.send_usb_req(b"\x01\x00",2, seqnum=usb_req.seqnum)
            pass
        return handled


'''
To run:
    on SERVER::
    $ python hid_device.py

    on CLIENT::
    sudo modprobe vhci-hcd
    sudo usbip -d list -r 127.0.0.1
    sudo usbip -d attach -r 127.0.0.1 -b 1-1.1
    lsusb -v -d 1337:1337
'''

usb_dev = CTAP2HIDevice()
usb_container = USBContainer()
usb_container.add_usb_device(usb_dev)  # Supports only one device!
usb_container.run()
