// Copyright 2015-2017 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 super::{WalletInfo, KeyPath};

use bigint::hash::H256;
use ethkey::{Address, Signature};
use hidapi;
use parking_lot::{Mutex, RwLock};

use std::cmp::min;
use std::fmt;
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;

const LEDGER_VID: u16 = 0x2c97;
const LEDGER_PIDS: [u16; 2] = [0x0000, 0x0001]; // Nano S and Blue
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

const APDU_TAG: u8 = 0x05;
const APDU_CLA: u8 = 0xe0;

#[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;
}

/// Hardware wallet error.
#[derive(Debug)]
pub enum Error {
	/// Ethereum wallet protocol error.
	Protocol(&'static str),
	/// Hidapi error.
	Usb(hidapi::HidError),
	/// Device with request key is not available.
	KeyNotFound,
	/// Signing has been cancelled by user.
	UserCancel,
}

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::KeyNotFound => write!(f, "Key not found"),
			Error::UserCancel => write!(f, "Operation has been cancelled"),
		}
	}
}

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

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

#[derive(Debug)]
struct Device {
	path: String,
	info: WalletInfo,
}

impl Manager {
	/// Create a new instance.
	pub fn new(hidapi: Arc<Mutex<hidapi::HidApi>>) -> Manager {
		Manager {
			usb: hidapi,
			devices: RwLock::new(Vec::new()),
			key_path: RwLock::new(KeyPath::Ethereum),
		}
	}

	/// Re-populate device list. Only those devices that have Ethereum app open will be added.
	pub fn update_devices(&self) -> Result<usize, Error> {
		let mut usb = self.usb.lock();
		usb.refresh_devices();
		let devices = usb.devices();
		let mut new_devices = Vec::new();
		let mut num_new_devices = 0;
		for device in devices {
			trace!("Checking device: {:?}", device);
			if device.vendor_id != LEDGER_VID || !LEDGER_PIDS.contains(&device.product_id) {
				continue;
			}
			match self.read_device_info(&usb, &device) {
				Ok(info) => {
					debug!("Found device: {:?}", info);
					if !self.devices.read().iter().any(|d| d.path == info.path) {
						num_new_devices += 1;
					}
					new_devices.push(info);

				}
				Err(e) => debug!("Error reading device info: {}", e),
			};
		}
		*self.devices.write() = new_devices;
		Ok(num_new_devices)
	}

	/// Select key derivation path for a known chain.
	pub fn set_key_path(&self, key_path: KeyPath) {
		*self.key_path.write() = key_path;
	}

	fn read_device_info(&self, usb: &hidapi::HidApi, dev_info: &hidapi::HidDeviceInfo) -> Result<Device, Error> {
		let mut handle = self.open_path(|| usb.open_path(&dev_info.path))?;
		let address = Self::read_wallet_address(&mut handle, *self.key_path.read())?;
		let manufacturer = dev_info.manufacturer_string.clone().unwrap_or("Unknown".to_owned());
		let name = dev_info.product_string.clone().unwrap_or("Unknown".to_owned());
		let serial = dev_info.serial_number.clone().unwrap_or("Unknown".to_owned());
		Ok(Device {
			path: dev_info.path.clone(),
			info: WalletInfo {
				name: name,
				manufacturer: manufacturer,
				serial: serial,
				address: address,
			},
		})
	}

	fn read_wallet_address(handle: &hidapi::HidDevice, key_path: KeyPath) -> Result<Address, 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"));
		}

		let (major, minor, patch) = (ver[1], ver[2], ver[3]);
		if major < 1 || (major == 1 && minor == 0 && patch < 3) {
			return Err(Error::Protocol("App version 1.0.3 is required."));
		}

		let eth_path = &ETH_DERIVATION_PATH_BE[..];
		let etc_path = &ETC_DERIVATION_PATH_BE[..];
		let derivation_path = match key_path {
			KeyPath::Ethereum => eth_path,
			KeyPath::EthereumClassic => etc_path,
		};
		let key_and_address = Self::send_apdu(handle, 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(address)
	}

	/// List connected wallets. This only returns wallets that are ready to be used.
	pub fn list_devices(&self) -> Vec<WalletInfo> {
		self.devices.read().iter().map(|d| d.info.clone()).collect()
	}

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

	/// Sign transaction data with wallet managing `address`.
	pub fn sign_transaction(&self, address: &Address, data: &[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))?;

		let eth_path = &ETH_DERIVATION_PATH_BE[..];
		let etc_path = &ETC_DERIVATION_PATH_BE[..];
		let derivation_path = match *self.key_path.read() {
			KeyPath::Ethereum => eth_path,
			KeyPath::EthereumClassic => etc_path,
		};
		const MAX_CHUNK_SIZE: usize = 255;
		let mut chunk: [u8; MAX_CHUNK_SIZE] = [0; MAX_CHUNK_SIZE];
		&mut chunk[0..derivation_path.len()].copy_from_slice(derivation_path);
		let mut dest_offset = derivation_path.len();
		let mut data_pos = 0;
		let mut result;
		loop {
			let p1 = if data_pos == 0 { 0x00 } else { 0x80 };
			let dest_left = MAX_CHUNK_SIZE - dest_offset;
			let chunk_data_size = min(dest_left, data.len() - data_pos);
			&mut chunk[dest_offset..][0..chunk_data_size].copy_from_slice(&data[data_pos..][0..chunk_data_size]);
			result = Self::send_apdu(&handle, commands::SIGN_ETH_TRANSACTION, p1, 0, &chunk[0..(dest_offset + chunk_data_size)])?;
			dest_offset = 0;
			data_pos += chunk_data_size;
			if data_pos == data.len() {
				break;
			}
		}

		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))
	}

	fn open_path<R, F>(&self, f: F) -> Result<R, Error>
		where F: Fn() -> Result<R, &'static str>
	{
		let mut err = Error::KeyNotFound;
		// Try to open device a few times.
		for _ in 0..10 {
			match f() {
				Ok(handle) => return Ok(handle),
				Err(e) => err = From::from(e),
			}
			::std::thread::sleep(Duration::from_millis(200));
		}
		Err(err)
	}

	fn send_apdu(handle: &hidapi::HidDevice, command: u8, p1: u8, p2: u8, data: &[u8]) -> Result<Vec<u8>, Error> {
		const HID_PACKET_SIZE: usize = 64 + HID_PREFIX_ZERO;
		let mut offset = 0;
		let mut chunk_index = 0;
		loop {
			let mut hid_chunk: [u8; HID_PACKET_SIZE] = [0; HID_PACKET_SIZE];
			let mut chunk_size = if chunk_index == 0 { 12 } else { 5 };
			let size = min(64 - chunk_size, data.len() - offset);
			{
				let chunk = &mut hid_chunk[HID_PREFIX_ZERO..];
				&mut chunk[0..5].copy_from_slice(&[0x01, 0x01, APDU_TAG, (chunk_index >> 8) as u8, (chunk_index & 0xff) as u8 ]);

				if chunk_index == 0 {
					let data_len = data.len() + 5;
					&mut 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 ]);
				}

				&mut chunk[chunk_size..chunk_size + size].copy_from_slice(&data[offset..offset + size]);
				offset += size;
				chunk_size += size;
			}
			trace!("writing {:?}", &hid_chunk[..]);
			let n = handle.write(&hid_chunk[..])?;
			if n < chunk_size {
				return Err(Error::Protocol("Write data size mismatch"));
			}
			if offset == data.len() {
				break;
			}
			chunk_index += 1;
		}

		// read response
		chunk_index = 0;
		let mut message_size = 0;
		let mut message = Vec::new();
		loop {
			let mut chunk: [u8; HID_PACKET_SIZE] = [0; HID_PACKET_SIZE];
			let chunk_size = handle.read(&mut chunk)?;
			trace!("read {:?}", &chunk[..]);
			if chunk_size < 5 || 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;
			}
			chunk_index += 1;
		}
		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!("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 Ethereum app is running.")),
			0x6faa => Err(Error::Protocol("You 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)
	}
}

#[test]
fn smoke() {
	use rustc_hex::FromHex;
	let hidapi = Arc::new(Mutex::new(hidapi::HidApi::new().unwrap()));
	let manager = Manager::new(hidapi.clone());
	manager.update_devices().unwrap();
	for d in &*manager.devices.read() {
		println!("Device: {:?}", d);
	}

	if let Some(address) = manager.list_devices().first().map(|d| d.address.clone()) {
		let tx = FromHex::from_hex("eb018504a817c80082520894a6ca2e6707f2cc189794a9dd459d5b05ed1bcd1c8703f26fcfb7a22480018080").unwrap();
		let signature = manager.sign_transaction(&address, &tx);
		println!("Got {:?}", signature);
		assert!(signature.is_ok());
		let large_tx = FromHex::from_hex("f8cb81968504e3b2920083024f279475b02a3c39710d6a3f2870d0d788299d48e790f180b8a4b61d27f6000000000000000000000000e1af840a5a1cb1efdf608a97aa632f4aa39ed199000000000000000000000000000000000000000000000000105ff43f46a9a800000000000000000000000000000000000000000000000000000000000000006000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000018080").unwrap();
		let signature = manager.sign_transaction(&address, &large_tx);
		println!("Got {:?}", signature);
		assert!(signature.is_ok());
		let huge_tx = FromHex::from_hex("[encoded data, 27608 characters]").unwrap();
		let signature = manager.sign_transaction(&address, &huge_tx);
		println!("Got {:?}", signature);
		assert!(signature.is_ok());
	}
}
