// Copyright 2015-2018 Parity Technologies (UK) Ltd.
// This file is part of Parity.

// Parity is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Parity is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Parity.  If not, see <http://www.gnu.org/licenses/>.

//! Ledger hardware wallet module. Supports Ledger Blue and Nano S.
//! See <https://github.com/LedgerHQ/blue-app-eth/blob/master/doc/ethapp.asc> for protocol details.

use std::cmp::min;
use std::str::FromStr;
use std::sync::{atomic, atomic::AtomicBool, Arc, Weak};
use std::time::{Duration, Instant};
use std::{fmt, thread};

use ethereum_types::{H256, Address};
use ethkey::Signature;
use hidapi;
use libusb;
use parking_lot::{Mutex, RwLock};
use semver::Version as FirmwareVersion;
use super::{WalletInfo, KeyPath, Device, DeviceDirection, Wallet, USB_DEVICE_CLASS_DEVICE, POLLING_DURATION};

const APDU_TAG: u8 = 0x05;
const APDU_CLA: u8 = 0xe0;
const APDU_PAYLOAD_HEADER_LEN: usize = 7;

const ETH_DERIVATION_PATH_BE: [u8; 17] = [4, 0x80, 0, 0, 44, 0x80, 0, 0, 60, 0x80, 0, 0, 0, 0, 0, 0, 0]; // 44'/60'/0'/0
const ETC_DERIVATION_PATH_BE: [u8; 21] = [5, 0x80, 0, 0, 44, 0x80, 0, 0, 60, 0x80, 0x02, 0x73, 0xd0, 0x80, 0, 0, 0, 0, 0, 0, 0]; // 44'/60'/160720'/0'/0

/// Ledger vendor ID
const LEDGER_VID: u16 = 0x2c97;
/// Ledger product IDs: [Nano S and Blue]
const LEDGER_PIDS: [u16; 2] = [0x0000, 0x0001];
const LEDGER_TRANSPORT_HEADER_LEN: usize = 5;

const MAX_CHUNK_SIZE: usize = 255;

const HID_PACKET_SIZE: usize = 64 + HID_PREFIX_ZERO;

#[cfg(windows)] const HID_PREFIX_ZERO: usize = 1;
#[cfg(not(windows))] const HID_PREFIX_ZERO: usize = 0;

mod commands {
	pub const GET_APP_CONFIGURATION: u8 = 0x06;
	pub const GET_ETH_PUBLIC_ADDRESS: u8 = 0x02;
	pub const SIGN_ETH_TRANSACTION: u8 = 0x04;
	pub const SIGN_ETH_PERSONAL_MESSAGE: u8 = 0x08;
}

/// Hardware wallet error.
#[derive(Debug)]
pub enum Error {
	/// Ethereum wallet protocol error.
	Protocol(&'static str),
	/// Hidapi error.
	Usb(hidapi::HidError),
	/// Libusb error
	LibUsb(libusb::Error),
	/// Device with request key is not available.
	KeyNotFound,
	/// Signing has been cancelled by user.
	UserCancel,
	/// Invalid device
	InvalidDevice,
	/// Impossible error
	Impossible,
	/// No device arrived
	NoDeviceArrived,
	/// No device left
	NoDeviceLeft,
}

impl fmt::Display for Error {
	fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
		match *self {
			Error::Protocol(ref s) => write!(f, "Ledger protocol error: {}", s),
			Error::Usb(ref e) => write!(f, "USB communication error: {}", e),
			Error::LibUsb(ref e) => write!(f, "LibUSB communication error: {}", e),
			Error::KeyNotFound => write!(f, "Key not found"),
			Error::UserCancel => write!(f, "Operation has been cancelled"),
			Error::InvalidDevice => write!(f, "Unsupported product was entered"),
			Error::Impossible => write!(f, "Placeholder error"),
			Error::NoDeviceArrived => write!(f, "No device arrived"),
			Error::NoDeviceLeft=> write!(f, "No device left"),
		}
	}
}

impl From<hidapi::HidError> for Error {
	fn from(err: hidapi::HidError) -> Self {
		Error::Usb(err)
	}
}

impl From<libusb::Error> for Error {
	fn from(err: libusb::Error) -> Self {
		Error::LibUsb(err)
	}
}

/// Ledger device manager.
pub (crate) struct Manager {
	usb: Arc<Mutex<hidapi::HidApi>>,
	devices: RwLock<Vec<Device>>,
	key_path: RwLock<KeyPath>,
}

impl Manager {
	/// Create a new instance.
	pub fn new(hidapi: Arc<Mutex<hidapi::HidApi>>, exiting: Arc<AtomicBool>) -> Result<Arc<Self>, libusb::Error> {
		let manager = Arc::new(Self {
			usb: hidapi,
			devices: RwLock::new(Vec::new()),
			key_path: RwLock::new(KeyPath::Ethereum),
		});

		let usb_context = Arc::new(libusb::Context::new()?);
		let m = manager.clone();

		// Subscribe to all Ledger devices
		// This means that we need to check that the given productID is supported
		// None => LIBUSB_HOTPLUG_MATCH_ANY, in other words that all are subscribed to
		// More info can be found: <http://libusb.sourceforge.net/api-1.0/group__hotplug.html#gae6c5f1add6cc754005549c7259dc35ea>
		usb_context.register_callback(
			Some(LEDGER_VID), None, Some(USB_DEVICE_CLASS_DEVICE),
			Box::new(EventHandler::new(Arc::downgrade(&manager)))).expect("usb_callback");

		// Ledger event handler thread
		thread::Builder::new()
			.spawn(move || {
				if let Err(e) = m.update_devices(DeviceDirection::Arrived) {
					debug!(target: "hw", "Ledger couldn't connect at startup, error: {}", e);
				}
				loop {
					usb_context.handle_events(Some(Duration::from_millis(500)))
						.unwrap_or_else(|e| debug!(target: "hw", "Ledger event handler error: {}", e));
					if exiting.load(atomic::Ordering::Acquire) {
						break;
					}
				}
			})
			.ok();

		Ok(manager)
	}

	// Transport Protocol:
	//		* Communication Channel Id		(2 bytes big endian )
	//		* Command Tag					(1 byte) 
	//		* Packet Sequence ID			(2 bytes big endian)
	//		* Payload						(Optional)
	//
	// Payload
	//		* APDU Total Length				(2 bytes big endian)
	//		* APDU_CLA						(1 byte)
	//		* APDU_INS						(1 byte)
	//		* APDU_P1						(1 byte)
	//		* APDU_P2						(1 byte)
	//		* APDU_LENGTH					(1 byte)
	//		* APDU_Payload					(Variable)
	// 
	fn write(handle: &hidapi::HidDevice, command: u8, p1: u8, p2: u8, data: &[u8]) -> Result<(), Error> {
		let data_len = data.len();
		let mut offset = 0;
		let mut sequence_number = 0;
		let mut hid_chunk = [0_u8; HID_PACKET_SIZE];
		
		while sequence_number == 0 || offset < data_len {
			let header = if sequence_number == 0 { LEDGER_TRANSPORT_HEADER_LEN + APDU_PAYLOAD_HEADER_LEN } else { LEDGER_TRANSPORT_HEADER_LEN };
			let size = min(64 - header, data_len - offset);
			{
				let chunk = &mut hid_chunk[HID_PREFIX_ZERO..];
				chunk[0..5].copy_from_slice(&[0x01, 0x01, APDU_TAG, (sequence_number >> 8) as u8, (sequence_number & 0xff) as u8 ]);

				if sequence_number == 0 {
					let data_len = data.len() + 5;
					chunk[5..12].copy_from_slice(&[(data_len >> 8) as u8, (data_len & 0xff) as u8, APDU_CLA, command, p1, p2, data.len() as u8]);
				}

				chunk[header..header + size].copy_from_slice(&data[offset..offset + size]);
			}
			trace!(target: "hw", "Ledger write {:?}", &hid_chunk[..]);
			let n = handle.write(&hid_chunk[..])?;
			if n < size + header {
				return Err(Error::Protocol("Write data size mismatch"));
			}
			offset += size;
			sequence_number += 1;
			if sequence_number >= 0xffff {
				return Err(Error::Protocol("Maximum sequence number reached"));
			}
		}
		Ok(())
	}
	
	// Transport Protocol:
	//		* Communication Channel Id		(2 bytes big endian )
	//		* Command Tag					(1 byte) 
	//		* Packet Sequence ID			(2 bytes big endian)
	//		* Payload						(Optional)
	//
	// Payload
	//		* APDU Total Length				(2 bytes big endian)
	//		* APDU_CLA						(1 byte)
	//		* APDU_INS						(1 byte)
	//		* APDU_P1						(1 byte)
	//		* APDU_P2						(1 byte)
	//		* APDU_LENGTH					(1 byte)
	//		* APDU_Payload					(Variable)
	// 
	fn read(handle: &hidapi::HidDevice) -> Result<Vec<u8>, Error> {
		let mut message_size = 0;
		let mut message = Vec::new();

		// terminate the loop if `sequence_number` reaches its max_value and report error
		for chunk_index in 0..=0xffff {
			let mut chunk: [u8; HID_PACKET_SIZE] = [0; HID_PACKET_SIZE];
			let chunk_size = handle.read(&mut chunk)?;
			trace!(target: "hw", "Ledger read {:?}", &chunk[..]);
			if chunk_size < LEDGER_TRANSPORT_HEADER_LEN || chunk[0] != 0x01 || chunk[1] != 0x01 || chunk[2] != APDU_TAG {
				return Err(Error::Protocol("Unexpected chunk header"));
			}
			let seq = (chunk[3] as usize) << 8 | (chunk[4] as usize);
			if seq != chunk_index {
				return Err(Error::Protocol("Unexpected chunk header"));
			}

			let mut offset = 5;
			if seq == 0 {
				// Read message size and status word.
				if chunk_size < 7 {
					return Err(Error::Protocol("Unexpected chunk header"));
				}
				message_size = (chunk[5] as usize) << 8 | (chunk[6] as usize);
				offset += 2;
			}
			message.extend_from_slice(&chunk[offset..chunk_size]);
			message.truncate(message_size);
			if message.len() == message_size {
				break;
			}
		}
		if message.len() < 2 {
			return Err(Error::Protocol("No status word"));
		}
		let status = (message[message.len() - 2] as usize) << 8 | (message[message.len() - 1] as usize);
		debug!(target: "hw", "Read status {:x}", status);
		match status {
			0x6700 => Err(Error::Protocol("Incorrect length")),
			0x6982 => Err(Error::Protocol("Security status not satisfied (Canceled by user)")),
			0x6a80 => Err(Error::Protocol("Invalid data")),
			0x6a82 => Err(Error::Protocol("File not found")),
			0x6a85 => Err(Error::UserCancel),
			0x6b00 => Err(Error::Protocol("Incorrect parameters")),
			0x6d00 => Err(Error::Protocol("Not implemented. Make sure the Ledger Ethereum Wallet app is running.")),
			0x6faa => Err(Error::Protocol("Your Ledger need to be unplugged")),
			0x6f00...0x6fff => Err(Error::Protocol("Internal error")),
			0x9000 => Ok(()),
			_ => Err(Error::Protocol("Unknown error")),

		}?;
		let new_len = message.len() - 2;
		message.truncate(new_len);
		Ok(message)
	}

	fn send_apdu(handle: &hidapi::HidDevice, command: u8, p1: u8, p2: u8, data: &[u8]) -> Result<Vec<u8>, Error> {
		Self::write(&handle, command, p1, p2, data)?;
		Self::read(&handle)
	}

	fn is_valid_ledger(device: &libusb::Device) -> Result<(), Error> {
		let desc = device.device_descriptor()?;
		let vendor_id = desc.vendor_id();
		let product_id = desc.product_id();

		if vendor_id == LEDGER_VID && LEDGER_PIDS.contains(&product_id) {
			Ok(())
		} else {
			Err(Error::InvalidDevice)
		}
	}

	fn get_firmware_version(handle: &hidapi::HidDevice) -> Result<FirmwareVersion, Error> {
		let ver = Self::send_apdu(&handle, commands::GET_APP_CONFIGURATION, 0, 0, &[])?;
		if ver.len() != 4 {
			return Err(Error::Protocol("Version packet size mismatch"));
		}
		Ok(FirmwareVersion::new(ver[1].into(), ver[2].into(), ver[3].into()))
	}

	fn get_derivation_path(&self) -> &[u8] {
		match *self.key_path.read() {
			KeyPath::Ethereum => &ETH_DERIVATION_PATH_BE,
			KeyPath::EthereumClassic => &ETC_DERIVATION_PATH_BE,
		}
	}
	
	fn signer_helper(&self, address: &Address, data: &[u8], command: u8) -> Result<Signature, Error> {
		let usb = self.usb.lock();
		let devices = self.devices.read();
		let device = devices.iter().find(|d| &d.info.address == address).ok_or(Error::KeyNotFound)?;
		let handle = self.open_path(|| usb.open_path(&device.path))?;

		// Signing personal messages are only support by Ledger firmware version 1.0.8 or newer
		if command == commands::SIGN_ETH_PERSONAL_MESSAGE {
			let version = Self::get_firmware_version(&handle)?;
			if version < FirmwareVersion::new(1, 0, 8) {
				return Err(Error::Protocol("Signing personal messages with Ledger requires version 1.0.8"));
			}
		}

		let mut chunk= [0_u8; MAX_CHUNK_SIZE];
		let derivation_path = self.get_derivation_path();

		// Copy the address of the key (only done once)
		chunk[0..derivation_path.len()].copy_from_slice(derivation_path);
		
		let key_length = derivation_path.len();
		let max_payload_size = MAX_CHUNK_SIZE - key_length;
		let data_len = data.len();
		
		let mut result = Vec::new();
		let mut offset = 0;
		
		while offset < data_len {
			let p1 = if offset == 0 { 0 } else { 0x80 };
			let take = min(max_payload_size, data_len - offset);

			// Fetch piece of data and copy it!
			{
				let (_key, d) = &mut chunk.split_at_mut(key_length);
				let (dst, _rem) = &mut d.split_at_mut(take);
				dst.copy_from_slice(&data[offset..(offset + take)]);
			}

			result = Self::send_apdu(&handle, command, p1, 0, &chunk[0..(key_length + take)])?;
			offset += take;
		}

		if result.len() != 65 {
			return Err(Error::Protocol("Signature packet size mismatch"));
		}
		let v = (result[0] + 1) % 2;
		let r = H256::from_slice(&result[1..33]);
		let s = H256::from_slice(&result[33..65]);
		Ok(Signature::from_rsv(&r, &s, v))
	}

	pub fn sign_message(&self, address: &Address, msg: &[u8]) -> Result<Signature, Error> {
		self.signer_helper(address, msg, commands::SIGN_ETH_PERSONAL_MESSAGE)
	}
}

// Try to connect to the device using polling in at most the time specified by the `timeout`
fn try_connect_polling(ledger: &Manager, timeout: &Duration, device_direction: DeviceDirection) -> bool {
	let start_time = Instant::now();
	while start_time.elapsed() <= *timeout {
		if let Ok(num_devices) = ledger.update_devices(device_direction) {
			trace!(target: "hw", "{} number of Ledger(s) {}", num_devices, device_direction);
			return true;
		}
	}
	false
}

impl <'a>Wallet<'a> for Manager {
	type Error = Error;
	type Transaction = &'a [u8];

	fn sign_transaction(&self, address: &Address, transaction: Self::Transaction) -> Result<Signature, Self::Error> {
		self.signer_helper(address, transaction, commands::SIGN_ETH_TRANSACTION)
	}
	
	fn set_key_path(&self, key_path: KeyPath) {
		*self.key_path.write() = key_path;
	}

	fn update_devices(&self, device_direction: DeviceDirection) -> Result<usize, Self::Error> {
		let mut usb = self.usb.lock();
		usb.refresh_devices();
		let devices = usb.devices();
		let num_prev_devices = self.devices.read().len();

		let detected_devices = devices.iter()
			.filter(|&d| d.vendor_id == LEDGER_VID && LEDGER_PIDS.contains(&d.product_id))
			.fold(Vec::new(), |mut v, d| {
				match self.read_device(&usb, &d) {
					Ok(info) => {
						trace!(target: "hw", "Found device: {:?}", info);
						v.push(info);
					}
					Err(e) => trace!(target: "hw", "Error reading device info: {}", e),
				};
				v
			});

		let num_curr_devices = detected_devices.len();
		*self.devices.write() = detected_devices;

		match device_direction {
			DeviceDirection::Arrived => {
				if num_curr_devices > num_prev_devices {
					Ok(num_curr_devices - num_prev_devices)
				} else {
					Err(Error::NoDeviceArrived)
				}
			}
			DeviceDirection::Left => {
				if num_prev_devices > num_curr_devices {
					Ok(num_prev_devices- num_curr_devices)
				} else {
					Err(Error::NoDeviceLeft)
				}
			}
		}
	}

	fn read_device(&self, usb: &hidapi::HidApi, dev_info: &hidapi::HidDeviceInfo) -> Result<Device, Self::Error> {
		let handle = self.open_path(|| usb.open_path(&dev_info.path))?;
		let manufacturer = dev_info.manufacturer_string.clone().unwrap_or_else(|| "Unknown".to_owned());
		let name = dev_info.product_string.clone().unwrap_or_else(|| "Unknown".to_owned());
		let serial = dev_info.serial_number.clone().unwrap_or_else(|| "Unknown".to_owned());
		match self.get_address(&handle) {
			Ok(Some(addr)) => {
				Ok(Device {
					path: dev_info.path.clone(),
					info: WalletInfo {
						name,
						manufacturer,
						serial,
						address: addr,
					},
				})
			}
			// This variant is not possible, but the trait forces this return type
			Ok(None) => Err(Error::Impossible),
			Err(e) => Err(e),
		}
	}

	fn list_devices(&self) -> Vec<WalletInfo> {
		self.devices.read().iter().map(|d| d.info.clone()).collect()
	}

	// Not used because it is not supported by Ledger
	fn list_locked_devices(&self) -> Vec<String> {
		vec![]
	}

	fn get_wallet(&self, address: &Address) -> Option<WalletInfo> {
		self.devices.read().iter().find(|d| &d.info.address == address).map(|d| d.info.clone())
	}

	fn get_address(&self, device: &hidapi::HidDevice) -> Result<Option<Address>, Self::Error> {
		let ledger_version = Self::get_firmware_version(&device)?;
		if ledger_version < FirmwareVersion::new(1, 0, 3) {
			return Err(Error::Protocol("Ledger version 1.0.3 is required"));
		}

		let derivation_path = self.get_derivation_path();

		let key_and_address = Self::send_apdu(device, commands::GET_ETH_PUBLIC_ADDRESS, 0, 0, derivation_path)?;
		if key_and_address.len() != 107 { // 1 + 65 PK + 1 + 40 Addr (ascii-hex)
			return Err(Error::Protocol("Key packet size mismatch"));
		}
		let address_string = ::std::str::from_utf8(&key_and_address[67..107])
			.map_err(|_| Error::Protocol("Invalid address string"))?;

		let address = Address::from_str(&address_string)
			.map_err(|_| Error::Protocol("Invalid address string"))?;

		Ok(Some(address))
	}

	fn open_path<R, F>(&self, f: F) -> Result<R, Self::Error>
		where F: Fn() -> Result<R, &'static str>
	{
		f().map_err(Into::into)
	}
}

/// Ledger event handler
/// A separate thread is handling incoming events
///
/// Note, that this run to completion and race-conditions can't occur but this can
/// therefore starve other events for being process with a spinlock or similar
struct EventHandler {
	ledger: Weak<Manager>,
}

impl EventHandler {
	/// Ledger event handler constructor
	fn new(ledger: Weak<Manager>) -> Self {
		Self { ledger }
	}
}

impl libusb::Hotplug for EventHandler {
	fn device_arrived(&mut self, device: libusb::Device) {
		debug!(target: "hw", "Ledger arrived");
		if let (Some(ledger), Ok(_)) = (self.ledger.upgrade(), Manager::is_valid_ledger(&device)) {
			if try_connect_polling(&ledger, &POLLING_DURATION, DeviceDirection::Arrived) != true {
				debug!(target: "hw", "No Ledger device was connected");
			}
		}
	}

	fn device_left(&mut self, device: libusb::Device) {
		debug!(target: "hw", "Ledger left");
		if let (Some(ledger), Ok(_)) = (self.ledger.upgrade(), Manager::is_valid_ledger(&device)) {
			if try_connect_polling(&ledger, &POLLING_DURATION, DeviceDirection::Left) != true {
				debug!(target: "hw", "No Ledger device was disconnected");
			}
		}
	}
}


#[cfg(test)]
mod tests {
	use rustc_hex::FromHex;
	use super::*;

	/// This test can't be run without an actual ledger device connected with the `Ledger Wallet Ethereum application` running
	#[test]
	#[ignore]
	fn sign_personal_message() {
		let manager = Manager::new(
			Arc::new(Mutex::new(hidapi::HidApi::new().expect("HidApi"))),
			Arc::new(AtomicBool::new(false))
		).expect("HardwareWalletManager");

		// Update device list
		manager.update_devices(DeviceDirection::Arrived).expect("No Ledger found, make sure you have a unlocked Ledger connected with the Ledger Wallet Ethereum running");

		// Fetch the ethereum address of a connected ledger device
		let address = manager.list_devices()
			.iter()
			.filter(|d| d.manufacturer == "Ledger".to_string())
			.nth(0)
			.map(|d| d.address.clone())
			.expect("No ledger device detected");

		// 44 bytes transaction
		let tx = FromHex::from_hex("eb018504a817c80082520894a6ca2e6707f2cc189794a9dd459d5b05ed1bcd1c8703f26fcfb7a22480018080").unwrap();
		let signature = manager.sign_transaction(&address, &tx);
		assert!(signature.is_ok());
	}

	/// This test can't be run without an actual ledger device connected with the `Ledger Wallet Ethereum application` running
	#[test]
	#[ignore]
	fn smoke() {
		let manager = Manager::new(
			Arc::new(Mutex::new(hidapi::HidApi::new().expect("HidApi"))),
			Arc::new(AtomicBool::new(false))
		).expect("HardwareWalletManager");

		// Update device list
		manager.update_devices(DeviceDirection::Arrived).expect("No Ledger found, make sure you have a unlocked Ledger connected with the Ledger Wallet Ethereum running");

		// Fetch the ethereum address of a connected ledger device
		let address = manager.list_devices()
			.iter()
			.filter(|d| d.manufacturer == "Ledger".to_string())
			.nth(0)
			.map(|d| d.address.clone())
			.expect("No ledger device detected");

		// 44 bytes transaction
		let tx = FromHex::from_hex("eb018504a817c80082520894a6ca2e6707f2cc189794a9dd459d5b05ed1bcd1c8703f26fcfb7a22480018080").unwrap();
		let signature = manager.sign_transaction(&address, &tx);
		println!("Got {:?}", signature);
		assert!(signature.is_ok());

		// 218 bytes transaction
		let large_tx = FromHex::from_hex("f86b028511cfc15d00825208940975ca9f986eee35f5cbba2d672ad9bc8d2a08448766c92c5cf830008026a0d2b0d401b543872d2a6a50de92455decbb868440321bf63a13b310c069e2ba5ba03c6d51bcb2e1653be86546b87f8a12ddb45b6d4e568420299b96f64c19701040f86b028511cfc15d00825208940975ca9f986eee35f5cbba2d672ad9bc8d2a08448766c92c5cf830008026a0d2b0d401b543872d2a6a50de92455decbb868440321bf63a13b310c069e2ba5ba03c6d51bcb2e1653be86546b87f8a12ddb45b6d4e568420299b96f64c19701040").unwrap();
		let signature = manager.sign_transaction(&address, &large_tx);
		println!("Got {:?}", signature);
		assert!(signature.is_ok());

		// 36206 bytes transaction (You need to confirm many transaction on your `Ledger` for this)
		let huge_tx = FromHex::from_hex("[encoded data, 36206 characters]
").unwrap();
		let signature = manager.sign_transaction(&address, &huge_tx);
		println!("Got {:?}", signature);
		assert!(signature.is_ok());
	}
}
