use std::{convert::TryFrom, str::FromStr};
use super::{errors::Error, wallet_helpers, KeyInfo, KeyStore};
use crate::shim::{address::Address, crypto::SignatureType};
use serde::{Deserialize, Serialize};
#[cfg(test)]
use {
crate::shim::crypto::Signature,
ahash::{HashMap, HashMapExt as _},
};
#[derive(Clone, PartialEq, Debug, Eq, Serialize, Deserialize)]
pub struct Key {
pub key_info: KeyInfo,
pub public_key: Vec<u8>,
pub address: Address,
}
impl TryFrom<KeyInfo> for Key {
type Error = crate::key_management::errors::Error;
fn try_from(key_info: KeyInfo) -> Result<Self, Self::Error> {
let public_key = wallet_helpers::to_public(*key_info.key_type(), key_info.private_key())?;
let address = wallet_helpers::new_address(*key_info.key_type(), &public_key)?;
Ok(Key {
key_info,
public_key,
address,
})
}
}
#[cfg(test)]
#[derive(Clone, PartialEq, Debug, Eq)]
pub struct Wallet {
keys: HashMap<Address, Key>,
keystore: KeyStore,
}
#[cfg(test)]
impl Wallet {
pub fn new(keystore: KeyStore) -> Self {
Wallet {
keys: HashMap::new(),
keystore,
}
}
pub fn new_from_keys(keystore: KeyStore, key_vec: impl IntoIterator<Item = Key>) -> Self {
let mut keys: HashMap<Address, Key> = HashMap::new();
for item in key_vec.into_iter() {
keys.insert(item.address, item);
}
Wallet { keys, keystore }
}
pub fn find_key(&mut self, addr: &Address) -> Result<Key, Error> {
if let Some(k) = self.keys.get(addr) {
return Ok(k.clone());
}
let key_string = format!("wallet-{addr}");
let key_info = match self.keystore.get(&key_string) {
Ok(k) => k,
Err(_) => {
self.keystore
.get(&format!("wallet-t{}", &addr.to_string()[1..]))?
}
};
let new_key = Key::try_from(key_info)?;
self.keys.insert(*addr, new_key.clone());
Ok(new_key)
}
pub fn sign(&mut self, addr: &Address, msg: &[u8]) -> Result<Signature, Error> {
let key = self.find_key(addr).map_err(|_| Error::KeyNotExists)?;
wallet_helpers::sign(*key.key_info.key_type(), key.key_info.private_key(), msg)
}
pub fn export(&mut self, addr: &Address) -> Result<KeyInfo, Error> {
let k = self.find_key(addr)?;
Ok(k.key_info)
}
pub fn import(&mut self, key_info: KeyInfo) -> Result<Address, Error> {
let k = Key::try_from(key_info)?;
let addr = format!("wallet-{}", k.address);
self.keystore.put(&addr, k.key_info)?;
Ok(k.address)
}
pub fn list_addrs(&self) -> Result<Vec<Address>, Error> {
list_addrs(&self.keystore)
}
pub fn get_default(&self) -> Result<Address, Error> {
let key_info = self.keystore.get("default")?;
let k = Key::try_from(key_info)?;
Ok(k.address)
}
pub fn set_default(&mut self, addr: Address) -> anyhow::Result<()> {
let addr_string = format!("wallet-{addr}");
let key_info = self.keystore.get(&addr_string)?;
if self.keystore.get("default").is_ok() {
self.keystore.remove("default")?; }
self.keystore.put("default", key_info)?;
Ok(())
}
pub fn generate_addr(&mut self, typ: SignatureType) -> anyhow::Result<Address> {
let key = generate_key(typ)?;
let addr = format!("wallet-{}", key.address);
self.keystore.put(&addr, key.key_info.clone())?;
self.keys.insert(key.address, key.clone());
let value = self.keystore.get("default");
if value.is_err() {
self.keystore
.put("default", key.key_info.clone())
.map_err(|err| Error::Other(err.to_string()))?;
}
Ok(key.address)
}
pub fn has_key(&mut self, addr: &Address) -> bool {
self.find_key(addr).is_ok()
}
}
pub fn get_default(keystore: &KeyStore) -> Result<Option<Address>, Error> {
if let Ok(key_info) = keystore.get("default") {
let k = Key::try_from(key_info)?;
Ok(Some(k.address))
} else {
Ok(None)
}
}
pub fn list_addrs(keystore: &KeyStore) -> Result<Vec<Address>, Error> {
let mut all = keystore.list();
all.sort();
let mut out = Vec::new();
for i in all {
if let Some(addr_str) = i.strip_prefix("wallet-") {
if let Ok(addr) = Address::from_str(addr_str) {
out.push(addr);
}
}
}
Ok(out)
}
pub fn find_key(addr: &Address, keystore: &KeyStore) -> Result<Key, Error> {
let key_string = format!("wallet-{addr}");
let key_info = keystore.get(&key_string)?;
let new_key = Key::try_from(key_info)?;
Ok(new_key)
}
pub fn remove_key(addr: &Address, keystore: &mut KeyStore) -> Result<(), Error> {
let key_string = format!("wallet-{addr}");
let deleted_keyinfo = keystore
.remove(&key_string)
.map_err(|_| Error::KeyNotExists)?;
if let Ok(default_keyinfo) = keystore.get("default") {
if default_keyinfo == deleted_keyinfo {
keystore
.remove("default")
.map_err(|_| Error::KeyNotExists)?;
}
}
println!("wallet {} deleted", addr);
Ok(())
}
pub fn try_find(addr: &Address, keystore: &mut KeyStore) -> Result<KeyInfo, Error> {
let key_string = format!("wallet-{addr}");
match keystore.get(&key_string) {
Ok(k) => Ok(k),
Err(_) => {
let mut new_addr = addr.to_string();
if new_addr.len() < 2 {
return Err(Error::Other(format!("Invalid addr {new_addr}")));
}
new_addr.replace_range(0..1, "t");
let key_string = format!("wallet-{new_addr}");
let key_info = match keystore.get(&key_string) {
Ok(k) => k,
#[allow(clippy::indexing_slicing)]
Err(_) => keystore.get(&format!("wallet-f{}", &new_addr[1..]))?,
};
Ok(key_info)
}
}
}
pub fn export_key_info(addr: &Address, keystore: &KeyStore) -> Result<KeyInfo, Error> {
let key = find_key(addr, keystore)?;
Ok(key.key_info)
}
pub fn generate_key(typ: SignatureType) -> Result<Key, Error> {
let private_key = wallet_helpers::generate(typ)?;
let key_info = KeyInfo::new(typ, private_key);
Key::try_from(key_info)
}
#[cfg(test)]
mod tests {
use crate::utils::encoding::blake2b_256;
use libsecp256k1::{Message as SecpMessage, SecretKey as SecpPrivate};
use super::*;
use crate::key_management::{generate, KeyStoreConfig};
fn construct_priv_keys() -> Vec<Key> {
let mut secp_keys = Vec::new();
let mut bls_keys = Vec::new();
for _ in 1..5 {
let secp_priv_key = generate(SignatureType::Secp256k1).unwrap();
let secp_key_info = KeyInfo::new(SignatureType::Secp256k1, secp_priv_key);
let secp_key = Key::try_from(secp_key_info).unwrap();
secp_keys.push(secp_key);
let bls_priv_key = generate(SignatureType::Bls).unwrap();
let bls_key_info = KeyInfo::new(SignatureType::Bls, bls_priv_key);
let bls_key = Key::try_from(bls_key_info).unwrap();
bls_keys.push(bls_key);
}
secp_keys.append(bls_keys.as_mut());
secp_keys
}
fn generate_wallet() -> Wallet {
let key_vec = construct_priv_keys();
Wallet::new_from_keys(KeyStore::new(KeyStoreConfig::Memory).unwrap(), key_vec)
}
#[test]
fn contains_key() {
let key_vec = construct_priv_keys();
let found_key = key_vec[0].clone();
let addr = key_vec[0].address;
let mut wallet =
Wallet::new_from_keys(KeyStore::new(KeyStoreConfig::Memory).unwrap(), key_vec);
assert_eq!(wallet.find_key(&addr).unwrap(), found_key);
assert!(wallet.has_key(&addr));
let new_priv_key = generate(SignatureType::Bls).unwrap();
let pub_key =
wallet_helpers::to_public(SignatureType::Bls, new_priv_key.as_slice()).unwrap();
let address = Address::new_bls(pub_key.as_slice()).unwrap();
assert!(!wallet.has_key(&address));
assert!(matches!(
wallet.find_key(&address).unwrap_err(),
Error::KeyInfo
));
assert!(!wallet.has_key(&address));
}
#[test]
fn sign() {
let key_vec = construct_priv_keys();
let priv_key_bytes = key_vec[2].key_info.private_key().clone();
let addr = key_vec[2].address;
let keystore = KeyStore::new(KeyStoreConfig::Memory).unwrap();
let mut wallet = Wallet::new_from_keys(keystore, key_vec);
let msg = [0u8; 64];
let msg_sig = wallet.sign(&addr, &msg).unwrap();
let msg_complete = blake2b_256(&msg);
let message = SecpMessage::parse(&msg_complete);
let priv_key = SecpPrivate::parse_slice(&priv_key_bytes).unwrap();
let (sig, recovery_id) = libsecp256k1::sign(&message, &priv_key);
let mut new_bytes = [0; 65];
new_bytes[..64].copy_from_slice(&sig.serialize());
new_bytes[64] = recovery_id.serialize();
let actual = Signature::new_secp256k1(new_bytes.to_vec());
assert_eq!(msg_sig, actual)
}
#[test]
fn import_export() {
let key_vec = construct_priv_keys();
let key = key_vec[0].clone();
let keystore = KeyStore::new(KeyStoreConfig::Memory).unwrap();
let mut wallet = Wallet::new_from_keys(keystore, key_vec);
let key_info = wallet.export(&key.address).unwrap();
assert_eq!(key_info, key.key_info);
let new_priv_key = generate(SignatureType::Secp256k1).unwrap();
let pub_key =
wallet_helpers::to_public(SignatureType::Secp256k1, new_priv_key.as_slice()).unwrap();
let test_addr = Address::new_secp256k1(pub_key.as_slice()).unwrap();
let key_info_err = wallet.export(&test_addr).unwrap_err();
assert!(matches!(key_info_err, Error::KeyInfo));
let test_key_info = KeyInfo::new(SignatureType::Secp256k1, new_priv_key);
assert!(wallet.import(test_key_info.clone()).is_ok());
let duplicate_error = wallet.import(test_key_info).unwrap_err();
assert!(matches!(duplicate_error, Error::KeyExists));
}
#[test]
fn list_addr() {
let key_vec = construct_priv_keys();
let mut addr_string_vec = Vec::new();
let mut key_store = KeyStore::new(KeyStoreConfig::Memory).unwrap();
for i in &key_vec {
addr_string_vec.push(i.address.to_string());
let addr_string = format!("wallet-{}", i.address);
key_store.put(&addr_string, i.key_info.clone()).unwrap();
}
addr_string_vec.sort();
let mut addr_vec = Vec::new();
for addr in addr_string_vec {
addr_vec.push(Address::from_str(addr.as_str()).unwrap())
}
let wallet = Wallet::new(key_store);
let test_addr_vec = wallet.list_addrs().unwrap();
assert_eq!(test_addr_vec, addr_vec);
}
#[test]
fn generate_new_key() {
let mut wallet = generate_wallet();
let addr = wallet.generate_addr(SignatureType::Bls).unwrap();
let key = wallet.keystore.get("default").unwrap();
assert_eq!(&SignatureType::Bls, key.key_type());
let address = format!("wallet-{addr}");
let key_info = wallet.keystore.get(&address).unwrap();
let key = wallet.keys.get(&addr).unwrap();
assert_eq!(key_info.key_type(), &SignatureType::Bls);
assert_eq!(key.address, addr);
}
#[test]
fn get_set_default() {
let key_store = KeyStore::new(KeyStoreConfig::Memory).unwrap();
let mut wallet = Wallet::new(key_store);
assert!(matches!(wallet.get_default().unwrap_err(), Error::KeyInfo));
let new_priv_key = generate(SignatureType::Secp256k1).unwrap();
let pub_key =
wallet_helpers::to_public(SignatureType::Secp256k1, new_priv_key.as_slice()).unwrap();
let test_addr = Address::new_secp256k1(pub_key.as_slice()).unwrap();
let key_info = KeyInfo::new(SignatureType::Secp256k1, new_priv_key);
let test_addr_string = format!("wallet-{test_addr}");
wallet.keystore.put(&test_addr_string, key_info).unwrap();
assert!(wallet.set_default(test_addr).is_ok());
assert_eq!(wallet.get_default().unwrap(), test_addr);
}
#[test]
fn secp_verify() {
let secp_priv_key = generate(SignatureType::Secp256k1).unwrap();
let secp_key_info = KeyInfo::new(SignatureType::Secp256k1, secp_priv_key);
let secp_key = Key::try_from(secp_key_info).unwrap();
let addr = secp_key.address;
let key_store = KeyStore::new(KeyStoreConfig::Memory).unwrap();
let mut wallet = Wallet::new_from_keys(key_store, vec![secp_key]);
let msg = [0u8; 64];
let sig = wallet.sign(&addr, &msg).unwrap();
sig.verify(&msg, &addr).unwrap();
let invalid_addr = wallet.generate_addr(SignatureType::Secp256k1).unwrap();
assert!(sig.verify(&msg, &invalid_addr).is_err())
}
#[test]
fn bls_verify_test() {
let bls_priv_key = generate(SignatureType::Bls).unwrap();
let bls_key_info = KeyInfo::new(SignatureType::Bls, bls_priv_key);
let bls_key = Key::try_from(bls_key_info).unwrap();
let addr = bls_key.address;
let key_store = KeyStore::new(KeyStoreConfig::Memory).unwrap();
let mut wallet = Wallet::new_from_keys(key_store, vec![bls_key]);
let msg = [0u8; 64];
let sig = wallet.sign(&addr, &msg).unwrap();
sig.verify(&msg, &addr).unwrap();
let invalid_addr = wallet.generate_addr(SignatureType::Bls).unwrap();
assert!(sig.verify(&msg, &invalid_addr).is_err())
}
}