# Copyright 2016 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""HID Transport for U2F.

This module imports the U2F HID Transport protocol as well as methods
for discovering devices implementing this protocol.
"""
from __future__ import absolute_import

import logging
import os
import struct
import time
import six

from . import errors, hid


def HidUsageSelector(device):
  if device['usage_page'] == 0xf1d0 and device['usage'] == 0x01:
    return True
  return False


def DiscoverLocalHIDU2FDevices(selector=HidUsageSelector):
  for d in hid.Enumerate():
    if selector(d):
      try:
        dev = hid.Open(d['path'])
        yield UsbHidTransport(dev)
      except OSError:
        # Insufficient permissions to access device
        pass


class UsbHidTransport(object):
  """Implements the U2FHID transport protocol.

  This class implements the U2FHID transport protocol from the
  FIDO U2F specs.  This protocol manages fragmenting longer messages
  over a short hid frame (usually 64 bytes).  It exposes an APDU
  channel through the MSG command as well as a series of other commands
  for configuring and interacting with the device.
  """

  U2FHID_PING = 0x81
  U2FHID_MSG = 0x83
  U2FHID_WINK = 0x88
  U2FHID_PROMPT = 0x87
  U2FHID_INIT = 0x86
  U2FHID_LOCK = 0x84
  U2FHID_ERROR = 0xbf
  U2FHID_SYNC = 0xbc

  CTAPHID_KEEPALIVE = 0xbb

  U2FHID_BROADCAST_CID = bytearray([0xff, 0xff, 0xff, 0xff])

  ERR_CHANNEL_BUSY = bytearray([0x06])


  class InitPacket(object):
    """Represent an initial U2FHID packet.

    Represent an initial U2FHID packet.  This packet contains
    metadata necessary to interpret the entire packet stream associated
    with a particular exchange (read or write).

    Attributes:
      packet_size: The size of the hid report (packet) used.  Usually 64.
      cid: The channel id for the connection to the device.
      size: The size of the entire message to be sent (including
          all continuation packets)
      payload: The portion of the message to put into the init packet.
          This must be smaller than packet_size - 7 (the overhead for
          an init packet).
    """

    def __init__(self, packet_size, cid, cmd, size, payload):
      self.packet_size = packet_size
      if len(cid) != 4 or cmd > 255 or size >= 2**16:
        raise errors.InvalidPacketError()
      if len(payload) > self.packet_size - 7:
        raise errors.InvalidPacketError()

      self.cid = cid  # byte array
      self.cmd = cmd  # number
      self.size = size  # number (full size of message)
      self.payload = payload  # byte array (for first packet)

    def ToWireFormat(self):
      """Serializes the packet."""
      ret = bytearray(64)
      ret[0:4] = self.cid
      ret[4] = self.cmd
      struct.pack_into('>H', ret, 5, self.size)
      ret[7:7 + len(self.payload)] = self.payload
      return list(six.iterbytes(bytes(ret)))

    @staticmethod
    def FromWireFormat(packet_size, data):
      """Derializes the packet.

      Deserializes the packet from wire format.

      Args:
        packet_size: The size of all packets (usually 64)
        data: List of ints or bytearray containing the data from the wire.

      Returns:
        InitPacket object for specified data

      Raises:
        InvalidPacketError: if the data isn't a valid InitPacket
      """
      ba = bytearray(data)
      if len(ba) != packet_size:
        raise errors.InvalidPacketError()
      cid = ba[0:4]
      cmd = ba[4]
      size = struct.unpack('>H', bytes(ba[5:7]))[0]
      payload = ba[7:7 + size]  # might truncate at packet_size
      return UsbHidTransport.InitPacket(packet_size, cid, cmd, size, payload)

  class ContPacket(object):
    """Represents a continutation U2FHID packet.

    Represents a continutation U2FHID packet.  These packets follow
    the intial packet and contains the remaining data in a particular
    message.

    Attributes:
      packet_size: The size of the hid report (packet) used.  Usually 64.
      cid: The channel id for the connection to the device.
      seq: The sequence number for this continuation packet.  The first
          continuation packet is 0 and it increases from there.
      payload:  The payload to put into this continuation packet.  This
          must be less than packet_size - 5 (the overhead of the
          continuation packet is 5).
    """

    def __init__(self, packet_size, cid, seq, payload):
      self.packet_size = packet_size
      self.cid = cid
      self.seq = seq
      self.payload = payload
      if len(payload) > self.packet_size - 5:
        raise errors.InvalidPacketError()
      if seq > 127:
        raise errors.InvalidPacketError()

    def ToWireFormat(self):
      """Serializes the packet."""
      ret = bytearray(self.packet_size)
      ret[0:4] = self.cid
      ret[4] = self.seq
      ret[5:5 + len(self.payload)] = self.payload
      return [int(x) for x in ret]

    @staticmethod
    def FromWireFormat(packet_size, data):
      """Derializes the packet.

      Deserializes the packet from wire format.

      Args:
        packet_size: The size of all packets (usually 64)
        data: List of ints or bytearray containing the data from the wire.

      Returns:
        InitPacket object for specified data

      Raises:
        InvalidPacketError: if the data isn't a valid ContPacket
      """
      ba = bytearray(data)
      if len(ba) != packet_size:
        raise errors.InvalidPacketError()
      cid = ba[0:4]
      seq = ba[4]
      # We don't know the size limit a priori here without seeing the init
      # packet, so truncation needs to be done in the higher level protocol
      # handling code, unlike the degenerate case of a 1 packet message in an
      # init packet, where the size is known.
      payload = ba[5:]
      return UsbHidTransport.ContPacket(packet_size, cid, seq, payload)

  def __init__(self, hid_device, read_timeout_secs=3.0):
    self.hid_device = hid_device

    in_size = hid_device.GetInReportDataLength()
    out_size = hid_device.GetOutReportDataLength()
    if in_size != out_size:
      raise errors.HardwareError(
          'unsupported device with different in/out packet sizes.')
    if in_size == 0:
      raise errors.HardwareError('unable to determine packet size')

    self.packet_size = in_size
    self.read_timeout_secs = read_timeout_secs
    self.logger = logging.getLogger('pyu2f.hidtransport')

    self.InternalInit()

  def SendMsgBytes(self, msg):
    r = self.InternalExchange(UsbHidTransport.U2FHID_MSG, msg)
    return r

  def SendBlink(self, length):
    return self.InternalExchange(UsbHidTransport.U2FHID_PROMPT,
                                 bytearray([length]))

  def SendWink(self):
    return self.InternalExchange(UsbHidTransport.U2FHID_WINK, bytearray([]))

  def SendPing(self, data):
    return self.InternalExchange(UsbHidTransport.U2FHID_PING, data)

  def InternalInit(self):
    """Initializes the device and obtains channel id."""
    self.cid = UsbHidTransport.U2FHID_BROADCAST_CID
    nonce = bytearray(os.urandom(8))
    r = self.InternalExchange(UsbHidTransport.U2FHID_INIT, nonce)
    if len(r) < 17:
      raise errors.HidError('unexpected init reply len')
    if r[0:8] != nonce:
      raise errors.HidError('nonce mismatch')
    self.cid = bytearray(r[8:12])

    self.u2fhid_version = r[12]
    self.device_version = tuple(r[13:16])
    self.capabilities = r[16]

  def InternalExchange(self, cmd, payload_in):
    """Sends and receives a message from the device."""
    # make a copy because we destroy it below
    self.logger.debug('payload: ' + str(list(payload_in)))
    payload = bytearray()
    payload[:] = payload_in
    for _ in range(2):
      self.InternalSend(cmd, payload)
      ret_cmd, ret_payload = self.InternalRecv()

      if ret_cmd == UsbHidTransport.U2FHID_ERROR:
        if ret_payload == UsbHidTransport.ERR_CHANNEL_BUSY:
          time.sleep(0.5)
          continue
        raise errors.HidError('Device error: %d' % int(ret_payload[0]))
      elif ret_cmd != cmd:
        raise errors.HidError('Command mismatch!')

      return ret_payload
    raise errors.HidError('Device Busy.  Please retry')

  def InternalSend(self, cmd, payload):
    """Sends a message to the device, including fragmenting it."""
    length_to_send = len(payload)

    max_payload = self.packet_size - 7
    first_frame = payload[0:max_payload]
    first_packet = UsbHidTransport.InitPacket(self.packet_size, self.cid, cmd,
                                              len(payload), first_frame)
    del payload[0:max_payload]
    length_to_send -= len(first_frame)
    self.InternalSendPacket(first_packet)

    seq = 0
    while length_to_send > 0:
      max_payload = self.packet_size - 5
      next_frame = payload[0:max_payload]
      del payload[0:max_payload]
      length_to_send -= len(next_frame)
      next_packet = UsbHidTransport.ContPacket(self.packet_size, self.cid, seq,
                                               next_frame)
      self.InternalSendPacket(next_packet)
      seq += 1

  def InternalSendPacket(self, packet):
    wire = packet.ToWireFormat()
    self.logger.debug('sending packet: ' + str(wire))
    self.hid_device.Write(wire)

  def InternalReadFrame(self):
        # TODO(gdasher): Figure out timeouts.  Today, this implementation
        # blocks forever at the HID level waiting for a response to a report.
        # This may not be reasonable behavior (though in practice in seems to be
        # OK on the set of devices and machines tested so far).
    frame = self.hid_device.Read()
    self.logger.debug('recv: ' + str(frame))
    return frame

  def InternalRecv(self):
    """Receives a message from the device, including defragmenting it."""
    first_read = self.InternalReadFrame()
    first_packet = UsbHidTransport.InitPacket.FromWireFormat(self.packet_size,
                                                             first_read)

    data = first_packet.payload
    to_read = first_packet.size - len(first_packet.payload)

    seq = 0
    while to_read > 0:
      next_read = self.InternalReadFrame()
      next_packet = UsbHidTransport.ContPacket.FromWireFormat(self.packet_size,
                                                              next_read)
      if self.cid != next_packet.cid:
        # Skip over packets that are for communication with other clients.
        # HID is broadcast, so we see potentially all communication from the
        # device.  For well-behaved devices, these should be BUSY messages
        # sent to other clients of the device because at this point we're
        # in mid-message transit.
        continue

      if seq != next_packet.seq:
        raise errors.HardwareError('Packets received out of order')

      # This packet for us at this point, so debit it against our
      # balance of bytes to read.
      to_read -= len(next_packet.payload)

      data.extend(next_packet.payload)
      seq += 1

    # truncate incomplete frames
    data = data[0:first_packet.size]
    return (first_packet.cmd, data)