extern crate libusb;
extern crate crypto;
extern crate rand;
#[macro_use]
extern crate bitflags;
use crypto::mac::Mac;
use rand::Rng;

use libusb::{DeviceHandle, Device};
pub use libusb::Context;

mod util;
use util::*;

mod handle;
use handle::*;

pub mod config;
pub use config::Config;

/// Get the vector of all YubiKeys in the libusb context.
pub fn get_yubikeys<'a>(ctx: &'a mut Context) -> Result<Vec<YubiKey<'a>>, Error> {

    let mut devices = Vec::new();
    for mut device in ctx.devices().unwrap().iter() {
        let descr = device.device_descriptor().unwrap();
        if descr.vendor_id() == 0x1050 && descr.product_id() == 0x407 {
            devices.push(YubiKey(device))
        }
    }
    Ok(devices)
}

/// A YubiKey device.
pub struct YubiKey<'a>(Device<'a>);

impl<'a> YubiKey<'a> {
    pub fn open(&mut self) -> Result<YubiKeyHandle<'a>, Error> {

        let mut h = try!(self.0.open());
        let config = self.0.config_descriptor(0).unwrap();
        let usb_int = config.interfaces().next().unwrap().descriptors().next().unwrap();

        if h.kernel_driver_active(0).unwrap() {
            h.detach_kernel_driver(0).unwrap();
        }

        h.set_active_configuration(1).unwrap_or(());
        h.claim_interface(usb_int.interface_number()).unwrap();

        Ok(YubiKeyHandle {
            h: h,
            version: None,
        })
    }
}

/// A handle to an open YubiKey.
pub struct YubiKeyHandle<'a> {
    h: DeviceHandle<'a>,
    version: Option<(Version, u8)>,
}

#[derive(Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Version {
    pub major: u8,
    pub minor: u8,
    pub build: u8,
}

#[derive(Clone, Copy, Debug)]
#[repr(u8)]
pub enum Command {
    Config = 0x01,
    Config2 = 0x03,
    Update1 = 0x04,
    Update2 = 0x05,
    Swap = 0x06,
    Ndef = 0x08,
    Ndef2 = 0x09,
    DeviceSerial = 0x10,
    DeviceConfig = 0x11,
    ScanMap = 0x12,
    Yk4Capabilities = 0x13,
    ChalOtp1 = 0x20,
    ChalOtp2 = 0x28,
    ChalHmac1 = 0x30,
    ChalHmac2 = 0x38,
}

/// A secret key for HMAC.
#[derive(Debug)]
pub struct HmacKey([u8; 20]);
impl Drop for HmacKey {
    fn drop(&mut self) {
        for i in self.0.iter_mut() {
            *i = 0;
        }
    }
}

/// A secret key for AES128 / OTP challenge-response.
#[derive(Debug)]
pub struct Aes128Key([u8; 16]);
impl Drop for Aes128Key {
    fn drop(&mut self) {
        for i in self.0.iter_mut() {
            *i = 0;
        }
    }
}

impl HmacKey {
    pub fn from_slice(s: &[u8]) -> Self {
        let mut key = HmacKey([0; 20]);
        (&mut key.0).clone_from_slice(s);
        key
    }

    pub fn generate<R:Rng>(mut rng: R) -> Self {
        let mut key = HmacKey([0; 20]);
        for i in key.0.iter_mut() {
            *i = rng.gen()
        }
        key
    }
}

impl Aes128Key {
    pub fn from_slice(s: &[u8]) -> Self {
        let mut key = Aes128Key([0; 16]);
        (&mut key.0).clone_from_slice(s);
        key
    }

    pub fn generate<R:Rng>(mut rng: R) -> Self {
        let mut key = Aes128Key([0; 16]);
        for i in key.0.iter_mut() {
            *i = rng.gen()
        }
        key
    }
}

#[derive(Debug)]
pub enum Error {
    WrongCRC,
    IO(std::io::Error),
    USB(libusb::Error),
    UnsupportedChallenge,
    ConfigNotWritten,
    CouldNotWrite
}
impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        use Error::*;
        match *self {
            WrongCRC => write!(f, "CRC failed"),
            IO(ref e) => e.fmt(f),
            USB(ref e) => e.fmt(f),
            UnsupportedChallenge => write!(f, "Unsupported challenge type"),
            ConfigNotWritten => write!(f, "Could not write configuration"),
            CouldNotWrite => write!(f, "Error writing to the yubikey"),
        }
    }
}

impl std::error::Error for Error {
    fn description(&self) -> &str {
        use Error::*;
        match *self {
            WrongCRC => "CRC failed",
            IO(ref e) => e.description(),
            USB(ref e) => e.description(),
            UnsupportedChallenge => "Unsupported challenge type",
            ConfigNotWritten => "Could not write configuration",
            CouldNotWrite => "Error writing to the yubikey",
        }
    }

    fn cause(&self) -> Option<&std::error::Error> {
        use Error::*;
        match *self {
            WrongCRC => None,
            IO(ref e) => Some(e),
            USB(ref e) => Some(e),
            UnsupportedChallenge => None,
            ConfigNotWritten => None,
            CouldNotWrite => None,
        }
    }
}


impl std::convert::From<libusb::Error> for Error {
    fn from(e: libusb::Error) -> Error {
        Error::USB(e)
    }
}

impl std::convert::From<std::io::Error> for Error {
    fn from(e: std::io::Error) -> Error {
        Error::IO(e)
    }
}

const VERSION_2_1: Version = Version {
    major: 2,
    minor: 1,
    build: 0,
};

const VERSION_2_2: Version = Version {
    major: 2,
    minor: 2,
    build: 0,
};

const VERSION_4: Version = Version {
    major: 4,
    minor: 0,
    build: 0,
};

/// Result of a challenge-response in OTP mode. The caller must check
/// that `(use_counter, session_counter)` (in alphabetical order) is
/// strictly larger than the last values seen.
#[repr(C)]
#[repr(packed)]
#[derive(Debug, Default)]
pub struct Otp {
    /// The private ID, XORed with the challenge.
    pub uid: [u8; 6],
    /// A counter incremented each time the YubiKey is powered up.
    pub use_counter: u16,
    /// A timestamp, encoded little-endian, starting from the time the YubiKey is powered up.
    pub timestamp: [u8; 3],
    /// A counter incremented with each response.
    pub session_counter: u8,
    /// A random number (not cryptographically strong).
    pub random_number: u16,
    /// CRC of the other fields.
    pub crc: u16,
}

/// Slot (1 or 2).
#[derive(Debug, Clone, Copy)]
pub enum Slot {
    Slot1,
    Slot2,
}

#[derive(Debug)]
pub struct Hmac([u8; 20]);
impl Drop for Hmac {
    fn drop(&mut self) {
        for i in self.0.iter_mut() {
            *i = 0;
        }
    }
}

impl std::ops::Deref for Hmac {
    type Target = [u8];
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Hmac {
    pub fn check(&self, key: &HmacKey, challenge: &[u8]) -> bool {
        &self.0[..] == hmac_sha1(key, challenge)
    }
}

#[derive(Debug)]
pub struct Aes128Block {
    block: [u8; 16],
}
impl Drop for Aes128Block {
    fn drop(&mut self) {
        for i in self.block.iter_mut() {
            *i = 0;
        }
    }
}

impl Aes128Block {
    /// Decrypts an AES block as returned by the YubiKey. The caller
    /// must check that the `uid` field is equal to the known private
    /// id, and that the `(use_counter, session_counter)` is strictly
    /// larger than the last value seen.
    pub fn check(&self, key: &Aes128Key, challenge: &[u8]) -> Result<Otp, Error> {

        let aes_dec = crypto::aessafe::AesSafe128Decryptor::new(&key.0);
        let mut tmp = Otp::default();
        {
            use crypto::symmetriccipher::BlockDecryptor;
            let mut tmp =
                unsafe { std::slice::from_raw_parts_mut(&mut tmp as *mut Otp as *mut u8, 16) };
            aes_dec.decrypt_block(&self.block, &mut tmp);

            if crc16(&tmp) != CRC_RESIDUAL_OK {
                return Err(Error::WrongCRC);
            }
        }

        for i in 0..6 {
            tmp.uid[i] ^= challenge[i]
        }

        Ok(tmp)
    }
}

impl<'a> YubiKeyHandle<'a> {
    /// Perform a challenge-response in HMAC-SHA1 mode. If `variable`
    /// is `true`, and the YubiKey is configured in variable length
    /// mode, challenges can be of any length up to 63 bytes. Else,
    /// challenges must be exactly 64 bytes long.
    pub fn challenge_response_hmac(&mut self,
                                   slot: Slot,
                                   variable: bool,
                                   challenge: &[u8])
                                   -> Result<Hmac, Error> {

        // Write the challenge.
        let mut challenge_ = if variable && challenge.last() == Some(&0) {
            [0xff; 64]
        } else {
            [0; 64]
        };
        (&mut challenge_[..challenge.len()]).copy_from_slice(challenge);
        let d = Frame::new(if let Slot::Slot1 = slot {
                               Command::ChalHmac1
                           } else {
                               Command::ChalHmac2
                           },
                           challenge_);
        let mut buf = [0; 8];
        try!(self.wait(|f| !f.contains(SLOT_WRITE_FLAG), &mut buf));

        if self.version.as_ref().unwrap().0 < VERSION_2_1 {
            return Err(Error::UnsupportedChallenge);
        }

        try!(self.write_frame(&d));

        // Read the response.
        let mut response = [0; 36];
        try!(self.read_response(&mut response));

        let mut hmac = Hmac([0; 20]);
        // Check response.
        if crc16(&response[..22]) != CRC_RESIDUAL_OK {
            return Err(Error::WrongCRC);
        }
        hmac.0.clone_from_slice(&response[..20]);
        Ok(hmac)
    }

    /// Perform a challenge-response in OTP mode. Returns the message
    /// from the YubiKey, decrypted.
    pub fn challenge_response_otp(&mut self,
                                  slot: Slot,
                                  challenge: &[u8; 6])
                                  -> Result<Aes128Block, Error> {

        // Write the challenge.
        let mut challenge_ = [0; 64];
        (&mut challenge_[..6]).copy_from_slice(challenge);
        let d = Frame::new(if let Slot::Slot1 = slot {
                               Command::ChalOtp1
                           } else {
                               Command::ChalOtp2
                           },
                           challenge_);
        let mut buf = [0; 8];
        try!(self.wait(|f| !f.contains(SLOT_WRITE_FLAG), &mut buf));

        if self.version.as_ref().unwrap().0 < VERSION_2_1 {
            return Err(Error::UnsupportedChallenge);
        }

        try!(self.write_frame(&d));

        // Read the response.
        let mut response = [0; 36];
        try!(self.read_response(&mut response));

        // Check response.
        if crc16(&response[..18]) != CRC_RESIDUAL_OK {
            return Err(Error::WrongCRC);
        }

        let mut block = Aes128Block { block: [0; 16] };
        block.block.copy_from_slice(&response[..16]);
        Ok(block)
    }

    /// Write the configuration to the yubikey. Currently only works
    /// with yubikeys >= 2.2 (this is checked).
    pub fn write_config(&mut self, command: Command, config: &mut Config) -> Result<(), Error> {


        let d = config.to_frame(command);
        let mut buf = [0; 8];
        try!(self.wait(|f| !f.contains(SLOT_WRITE_FLAG), &mut buf));

        // We should do a more careful check of the version number.
        assert!(self.version.as_ref().unwrap().0 >= VERSION_2_2);

        let old_pgm = self.version.as_ref().unwrap().1;

        try!(self.write_frame(&d));
        try!(self.wait(|f| !f.contains(SLOT_WRITE_FLAG), &mut buf));

        let pgm = self.version.as_ref().unwrap().1;

        if old_pgm == pgm {
            Err(Error::ConfigNotWritten)
        } else {
            Ok(())
        }
    }

    /// Read the capabilities of a yubikey >= 4.0. Not sure what to do with it for now.
    fn read_capabilities(&mut self) -> Result<(), Error> {
        let mut buf = [0; 8];
        try!(self.wait(|f| !f.contains(SLOT_WRITE_FLAG), &mut buf));

        assert!(self.version.as_ref().unwrap().0 >= VERSION_4);

        let d = Frame::new(Command::Yk4Capabilities, [0; 64]);
        try!(self.write_frame(&d));
        println!("written");
        let mut response = [0; 64];

        self.read_response(&mut response).unwrap();

        let r_len = response[0] as usize;

        assert_eq!(crc16(&response[..r_len + 3]), CRC_RESIDUAL_OK);

        let r = &response[1..r_len + 1];
        println!("{:?}", r);
        let mut i = 0;
        while i < r.len() {

            let t = Capabilities::from_bits(r[i]).unwrap();
            let l = r[i + 1] as usize;
            let v = &r[i + 2..i + 2 + l];

            println!("{:?} {:?}", t, v);

            i += 2 + l
        }
        Ok(())
    }
}

bitflags! {
    flags Capabilities: u8 {
        const OTP = 0x1,
        const U2F = 0x2,
        const CCID = 0x4,
    }
}



const SHA1_DIGEST_SIZE: usize = 20;
fn hmac_sha1(key: &HmacKey, data: &[u8]) -> [u8; SHA1_DIGEST_SIZE] {
    let digest = crypto::sha1::Sha1::new();
    let mut hmac = crypto::hmac::Hmac::new(digest, &key.0);
    hmac.input(data);

    let mut code = [0; SHA1_DIGEST_SIZE];
    hmac.raw_result(&mut code);

    code
}