rustwallet/src/wallet/wallet_impl.rs

extern crate hex;

use anyhow::Result;
use core::fmt::Write;
use primitive_types::{H160, H256};
use secp256k1::rand::rngs::OsRng;
use secp256k1::{Message, PublicKey, Secp256k1, SecretKey};
use serde::{Deserialize, Serialize};
use shamir_secret_sharing::num_bigint::BigInt;
use shamir_secret_sharing::ShamirSecretSharing as SSS;
use std::str;
use std::str::FromStr;
use tiny_keccak::keccak256;

pub fn generate_keypair() -> (SecretKey, PublicKey) {
    let secp = Secp256k1::new();
    // let mut rng = OsRng::new().expect("OsRng");
    secp.generate_keypair(&mut OsRng)
}

pub fn public_key_address(public_key: &PublicKey) -> H160 {
    let public_key = public_key.serialize_uncompressed();
    debug_assert_eq!(public_key[0], 0x04);
    let hash = keccak256(&public_key[1..]);

    H160::from_slice(&hash[12..])
}

pub fn get_public_key(secret_key: &SecretKey) -> PublicKey {
    let secp = Secp256k1::new();
    PublicKey::from_secret_key(&secp, secret_key)
}

pub fn get_sss(msg: &str) -> SSS {
    let k_hash = keccak256(msg.as_bytes());
    let s = hex::encode(&k_hash);
    let pb = BigInt::parse_bytes(&s.as_bytes(), 16).unwrap();
    SSS {
        threshold: 2,
        share_amount: 3,
        prime: pb,
    }
}

pub fn generate_sss_keypair(msg: &str, skey: &str) -> Vec<String> {
    let secret = BigInt::parse_bytes(&skey.as_bytes(), 16).unwrap();
    let sss = get_sss(msg);
    let shares = sss.split(secret.clone());
    let mut shares_str: Vec<String> = Vec::new();
    for (_i, v) in &shares {
        shares_str.push(v.to_str_radix(16));
    }
    shares_str
}

pub fn hash_message<S>(message: S) -> H256
where
    S: AsRef<[u8]>,
{
    let message = message.as_ref();

    let mut eth_message = format!("\x19Ethereum Signed Message:\n{}", message.len()).into_bytes();
    eth_message.extend_from_slice(message);

    keccak256(&eth_message).into()
}

#[derive(Serialize, Deserialize, Debug)]
pub struct Wallet {
    pub msg_key: String,
    pub master_key: String,
    pub second_key: Option<String>,
    pub backup_key: Option<String>,
}

impl Wallet {
    pub fn new(msg: &str) -> Self {
        let (secret_key, _pub_key) = generate_keypair();
        let s = hex::encode(&secret_key.secret_bytes());
        let shares_str = generate_sss_keypair(msg, &s);
        // println!("secret key: {:?}", secret_key);
        // println!("{:?}", s);
        let second_key = shares_str.get(1).map(String::clone);
        let backup_key = shares_str.get(2).map(String::clone);
        Wallet {
            msg_key: msg.to_string(),
            master_key: shares_str.get(0).unwrap().to_string(),
            second_key: second_key,
            backup_key: backup_key,
        }
    }

    pub fn reset_wallet(&self) -> Self {
        let secret_key = self.get_secret_key();
        let s = hex::encode(&secret_key.secret_bytes());
        let shares_str = generate_sss_keypair(&self.msg_key, &s);
        let second_key = shares_str.get(1).map(String::clone);
        let backup_key = shares_str.get(2).map(String::clone);

        Wallet {
            msg_key: self.msg_key.clone(),
            master_key: shares_str.get(0).unwrap().to_string(),
            second_key: second_key,
            backup_key: backup_key,
        }
    }

    pub fn get_secret_key(&self) -> SecretKey {
        let key_str_0: &str = &self.master_key;
        let key0 = BigInt::parse_bytes(&key_str_0.as_bytes(), 16).unwrap();
        let kp0: (usize, BigInt) = (1, key0);
        let (i, key_str_1) = match (&self.second_key, &self.backup_key) {
            (Some(val), _) => (2, val),
            (_, Some(val)) => (3, val),
            _ => {
                panic!("error generate key")
            }
        };
        let key1 = BigInt::parse_bytes(&key_str_1.as_bytes(), 16).unwrap();
        let kp1 = (i, key1);
        let _tmp = vec![kp0, kp1];
        let sss = get_sss(&self.msg_key);
        let secret_b = sss.recover(&_tmp);
        let mut s_key_str = secret_b.to_str_radix(16);
        if s_key_str.len() < 64 {
            s_key_str += "0";
        }
        SecretKey::from_str(&s_key_str).expect("32 bytes, within curve order")
    }

    pub fn get_public_key(&self) -> PublicKey {
        let s_key = self.get_secret_key();
        get_public_key(&s_key)
    }

    pub fn generate_sec_key(&self) -> String {
        let secret_key = self.get_secret_key();
        let s = hex::encode(&secret_key.secret_bytes());
        s
    }

    pub fn get_address(&self) -> H160 {
        let public_key = self.get_public_key();
        public_key_address(&public_key)
    }

    pub fn sign<S>(&self, msg: S) -> Result<String>
    where
        S: AsRef<[u8]>,
    {
        let secp = Secp256k1::new();
        let secret_key = self.get_secret_key();
        let message = msg.as_ref();
        let message_hash = hash_message(message.as_ref());
        let message_to_hash = Message::from_slice(message_hash.as_ref()).unwrap();
        let (recovery_id, signature) = secp
            .sign_ecdsa_recoverable(&message_to_hash, &secret_key)
            .serialize_compact();

        let mut s = hex::encode(signature);
        let standard_v = recovery_id.to_i32() as u64 + 27;
        let rv: u8 = standard_v
            .try_into()
            .expect("signature recovery in electrum notation always fits in a u8");
        write!(s, "{:02x}", rv).unwrap();
        Ok(s)
    }
    pub fn sign_for_tran<S>(&self, msg: S) -> Result<(String, i32)>
    where
        S: AsRef<[u8]>,
    {
        let secp = Secp256k1::new();
        let secret_key = self.get_secret_key();
        let hex_str = match hex::decode(msg) {
            Ok(v) => v,
            Err(e) => panic!("error decode hex str: {}", e),
        };
        let message_to_hash = Message::from_slice(&hex_str).unwrap();
        let (_recovery_id, signature) = secp
            .sign_ecdsa_recoverable(&message_to_hash, &secret_key)
            .serialize_compact();
        let s = hex::encode(signature);
        let recid = _recovery_id.to_i32();
        Ok((s, recid))
    }
}