# Run with a single argument: a /dev/hidrawX path.
# If you don't have udev setup to allow access to U2F tokens, you may need to
# chown the device to your user before running this script.
# If you don't know which hidraw to use, try removing and reinserting your
# token. Then the device with the most recent ctime is the one you want.
#
# Once running, press the token's button twice. The first press will trigger a
# registration, the second an authentication.

import sys
import struct
import time
import hashlib
import os

if len(sys.argv) < 2:
    print('Please give /dev/hidrawX path as first argument')

dev = file(sys.argv[1], 'rb+')

def zeroPad(requestedLen, x):
    return x + ('\x00' * max(requestedLen-len(x), 0))

chanID = ''

def write(usbCmd, payload):
  prefix = '\x00' + chanID

  i = -1
  while len(payload) > 0:
      if i < 0:
          packet = prefix + struct.pack('>BH', usbCmd, len(payload))
      else:
          packet = prefix + struct.pack('B', i)

      todo = min(len(payload), 65 - len(packet))
      packet += payload[:todo]
      payload = payload[todo:]
      dev.write(zeroPad(65, packet))
      dev.flush()
      #print('> ' + packet.encode('hex'))
      i += 1

def read():
    packet = dev.read(64)
    #print('< ' + packet.encode('hex'))
    packet = packet[4:]
    (usbCmd, length) = struct.unpack('>BH', packet[:3])
    packet = packet[3:]
    reply = packet[:min(length, len(packet))]

    while length > len(reply):
        packet = dev.read(64)
        packet = packet[5:]
        reply += packet[:min(length - len(reply), len(packet))]

    return usbCmd, reply

def transact(cmd, p1, p2, body):
    payload = struct.pack('BBBB', 0, cmd, p1, p2)
    if len(body) != 0:
        payload += struct.pack('>BH', 0, len(body))
    payload += body
    if len(body) != 0:
        payload += '\x00' * 3
    else:
        payload += '\x00' * 2
    write(0x83, payload)
    (usbCmd, reply) = read()
    return struct.unpack('>H', reply[-2:])[0], reply[:-2]

# Request a channel from the device.
dev.write(zeroPad(65, '00ffffffff860008cbf160f6e427b502'.decode('hex')))
reply = dev.read(64)
idOffset = 4 + 1 + 2 + 8
chanID = reply[idOffset:idOffset+4]

while True:
    challenge_hash = hashlib.sha256('challenge').digest()
    appid_hash = hashlib.sha256('https://example.com').digest()
    status, reply = transact(1, 3, 0, challenge_hash + appid_hash)
    if status == 0x6985:
        time.sleep(0.5)
        continue
    if status != 0x9000:
        print hex(status)
        os.exit(1)
    print reply.encode('hex')
    print 'Public key: ' + reply[1:66].encode('hex')
    key_handle_length = ord(reply[66])
    key_handle = reply[67:67 + key_handle_length]
    print 'Key handle: ' + key_handle.encode('hex')
    reply = reply[67 + key_handle_length:]
    # This is a fragile way to get an certificate length; just to keep
    # things short.
    cert_len = struct.unpack('>H', reply[2:4])[0]
    print '-----BEGIN CERTIFICATE-----'
    print reply[:4+cert_len].encode('base64'),
    print '-----END CERTIFICATE-----'
    print 'Signature: ' + reply[4+cert_len:].encode('hex')
    break

while True:
    challenge_hash = hashlib.sha256('challenge').digest()
    appid_hash = hashlib.sha256('https://example.com').digest()
    status, reply = transact(2, 3, 0, challenge_hash + appid_hash +
                                      chr(len(key_handle)) + key_handle)
    if status == 0x6985:
        time.sleep(0.5)
        continue
    if status != 0x9000:
        print hex(status)
        os.exit(1)
    (flags, counter) = struct.unpack('>BI', reply[:5])
    if flags & 1:
        print 'User-presence tested'
    print 'Counter: %d' % counter
    print 'Signature: ' + reply[5:].encode('hex')
    break