from fido2.hid import CtapHidDevice
from fido2.client import Fido2Client, UserInteraction, WindowsClient
from getpass import getpass
import sys, os ,ctypes, time, cmd, json

# Handle user interaction
class CliInteraction(UserInteraction):
    def prompt_up(self):
        print("\nTouch your authenticator device now...\n")

    def request_pin(self, permissions, rd_id):
        return getpass("Enter PIN: ")

    def request_uv(self, permissions, rd_id):
        print("User Verification required.")
        return True

def setup_windows():
    if WindowsClient.is_available() and not ctypes.windll.shell32.IsUserAnAdmin():
        # Use the Windows WebAuthn API if available, and we're not running as admin
        client = WindowsClient("https://storingsecrets.fido2")
        if "hmac-secret" in client.info.extensions:
            return client
        else:
            return None
            
def setup_pcsc():
    elevate()
    
    # Import PCSC
    try:
        from fido2.pcsc import CtapPcscDevice
    except ImportError:
        CtapPcscDevice = None

    # List devices
    def enumerate_devices():
        for dev in CtapHidDevice.list_devices():
            yield dev
        if CtapPcscDevice:
            for dev in CtapPcscDevice.list_devices():
                yield dev


    # Locate first device, prefer ctap
    for dev in enumerate_devices():
        client = Fido2Client(dev, "https://storingsecrets.fido2", user_interaction=CliInteraction())
        if "hmac-secret" in client.info.extensions:
            return client
    else:
        return None

def elevate():
    if ctypes.windll.shell32.IsUserAnAdmin():
        pass
    else:
        print('Windows API not available for hmac-secret, elevating program to use direct CTAP. Please follow instructions in new window')
        ctypes.windll.shell32.ShellExecuteW(None, "runas", sys.executable, " ".join(sys.argv), None, 1)
        exit(0)
    
def challenge(bytes=32):
    return os.urandom(bytes)
    
def setup():
    client = setup_windows()
    if client:
        return client
    
    client = setup_pcsc()
    if client:
        return client
        
    return None
    
def load_keys():
    try:
        with open('keys.json', 'r') as file:
            keys = json.load(file)['fido_keys']
        return keys
    except FileNotFoundError:
        with open('keys.json', 'w') as file:
            json.dump({'fido_keys': []}, file)
        return []
    
def append_key(user):
    keys = load_keys()
    keys.append(user)
    
    writethis = {}
    writethis['fido_keys'] = keys
    
    with open('keys.json', 'w') as file:
        json.dump(writethis, file)
    
class RootShell(cmd.Cmd):
    intro = f'Welcome {os.getlogin()} to my own shell for configuration.    Type help or ? to list nodes\n'
    prompt = '(root) '

    # Shell settings
    def default(self, arg):
        """ Return friendly string instead of bad syntax """
        print(f'Node does not exist, type help or ? to list commands.')

    def emptyline(self):
        """ Return empty """
        pass

    # Program functions
    def do_list(self, arg):
        """
        List available keys: LIST
        """
        keys = load_keys()
        
        if len(keys) == 0:
            print('No available keys')
            return None
        
        print("{:<8} {:<30}".format('Nr', 'Short name'))
        for n, v in enumerate(keys):
            print("{:<8} {:<30}".format(n, v['name']))

    def do_create(self, arg):
        """
        Create new key: CREATE
        """
        usage = input('What is the purpose of this key: ')
        
        user = {"id": challenge(8), "name": str(usage)}
        
        result = client.make_credential(
        {
            "rp": rp,
            "user": user,
            "challenge": challenge(),
            "pubKeyCredParams": [{"type": "public-key", "alg": -7}],
            "extensions": {"hmacCreateSecret": True},
        },)
        
        if not result.extension_results.get("hmacCreateSecret"):
            print("Failed to create credential with HmacSecret")
            return None
        
        credential = result.attestation_object.auth_data.credential_data
        
        append_key({'kid': credential.credential_id.hex(), 'name': usage, 'salt': challenge().hex()})
                
        print("New credential stored")
        print(f"KID (HEX): {credential.credential_id.hex()}")
        
        
    def do_sign(self, arg):
        """
        Perform key exchange for shared key: READ
        """
        self.do_list(None)
        usage = input('Which key should we use (Nr): ')
        
        key = load_keys()[int(usage)]
        
        allow_list = [{"type": "public-key", "id": bytes.fromhex(key['kid'])}]

        # Authenticate the credential
        result = client.get_assertion(
            {
                "rpId": rp["id"],
                "challenge": challenge(),
                "allowCredentials": allow_list,
                "userVerification": 'required',
                "extensions": {"hmacGetSecret": {"salt1": bytes.fromhex(key['salt'])}},
            },
        ).get_response(0)
        
        output1 = result.extension_results["hmacGetSecret"]["output1"]
        
        print(f"Authenticated secret (HEX): {output1.hex()}")
        
    def do_exit(self, arg):
        """
        Exit and close session: EXIT
        """
        print('Goodbye')
        return True


if __name__ == '__main__':
    client = setup()
    
    rp = {"id": "storingsecrets.fido2", "name": "localhost"}
    
    try:
        RootShell().cmdloop()
    except KeyboardInterrupt:
        print('\nGoodbye')
        exit(0)
    #except:
    #    print('Something went wrong')