Airgapped hardware wallet based on STM32F469-Discovery board using SD card to pass transaction data

Dependencies:   mbed QSPI_DISCO_F469NI BSP_DISCO_F469NI

main.cpp

Committer:
stepansnigirev
Date:
2019-07-20
Revision:
0:f43431023689
Child:
1:b88ef7630eb3

File content as of revision 0:f43431023689:

#include "mbed.h"
#include "helpers.h"

// bitcoin lib
#include "Bitcoin.h"
#include "PSBT.h"

/*********** forward declarations **************/

// generates a new mnemonic
string generateNewMnemonic();

// generates hd keys from the mnemonic
void initKeys(const string mnemonic, const string password = "");

/****************** GUI classes ****************/

Label titleLbl;
Label dataLbl;
QR qr;

Button btn;
Button printBtn;
Label lbl;

/***************** GUI functions ***************/

// if mnemonic is not present we can generate or recover it
void showInitScreen();
// if mnemonic is there we go directly to the main menu
void showMenu();
// handy function to display information to the user
// with a title, a message and OK button that goes to the menu
void showMessage(const string title, const string message);
static lv_res_t toMenuCallback(lv_obj_t * btn);
static lv_res_t showAddressesCallback(lv_obj_t * btn);
static lv_res_t saveXpubCallback(lv_obj_t * btn);
static lv_res_t signPSBTCallback(lv_obj_t * btn);
static lv_res_t showMnemonicCallback(lv_obj_t * btn);
static lv_res_t wipeCallback(lv_obj_t * btn);
static lv_res_t newMnemonicCallback(lv_obj_t * btn);
static lv_res_t enterMnemonicCallback(lv_obj_t * btn);
void showAddress(unsigned int child_index, bool change);

/***************** bitcoin keys ***************/

HDPrivateKey root; // root private key
HDPrivateKey account; // account master private key
HDPublicKey xpub; // account master public key
bool change = false; // internal or external address
unsigned int child_index = 0; // current child index

PSBT psbt; // psbt transaction we will be signing

/******************* Main part *****************/

int main(){
    init();

    string mnemonic = loadMnemonic();
    if(mnemonic.length() == 0){
        showInitScreen();
    }else{
        initKeys(mnemonic);
        showMenu();
    }

    while(1){
        gui.update();
    }
} 

/****************** GUI stuff *****************/
void showInitScreen(){
    gui.clear();
    titleLbl = Label("Let's set it up!");
    titleLbl.size(gui.width(), 20);
    titleLbl.position(0, 40);
    titleLbl.alignText(ALIGN_TEXT_CENTER);

    Button btn(newMnemonicCallback, "Generate new mnemonic");
    btn.size(gui.width()-100, 80);
    btn.position(0, 200);
    btn.align(ALIGN_CENTER);

    Button btn2(enterMnemonicCallback, "Enter existing mnemonic");
    btn2.size(gui.width()-100, 80);
    btn2.position(0, 300);
    btn2.align(ALIGN_CENTER);
}

void showMenu(){
    gui.clear();
    titleLbl = Label("What do you want to do?");
    titleLbl.size(gui.width(), 20);
    titleLbl.position(0, 40);
    titleLbl.alignText(ALIGN_TEXT_CENTER);

    Button btn(showAddressesCallback, "Show addresses");
    btn.size(gui.width()-100, 80);
    btn.position(0, 100);
    btn.align(ALIGN_CENTER);

    Button btn2(saveXpubCallback, "Export xpub");
    btn2.size(gui.width()-100, 80);
    btn2.position(0, 300);
    btn2.align(ALIGN_CENTER);

    Button btn3(signPSBTCallback, "Sign PSBT transaction");
    btn3.size(gui.width()-100, 80);
    btn3.position(0, 400);
    btn3.align(ALIGN_CENTER);

    Button btn4(wipeCallback, "Wipe device");
    btn4.size(gui.width()-100, 80);
    btn4.position(0, 600);
    btn4.align(ALIGN_CENTER);

    Button btn5(showMnemonicCallback, "Show mnemonic");
    btn5.size(gui.width()-100, 80);
    btn5.position(0, 700);
    btn5.align(ALIGN_CENTER);
}

void showMessage(const string title, const string message){
    gui.clear();
    titleLbl = Label(title);
    titleLbl.size(gui.width(), 20);
    titleLbl.position(0, 40);
    titleLbl.alignText(ALIGN_TEXT_CENTER);

    dataLbl = Label(message);
    dataLbl.size(gui.width()-100, 100);
    dataLbl.position(50, 300);
    dataLbl.alignText(ALIGN_TEXT_CENTER);

    Button btn(toMenuCallback, "OK");
    btn.size(gui.width()-100, 80);
    btn.position(0, gui.height()-100);
    btn.align(ALIGN_CENTER);
}

static lv_res_t nextCallback(lv_obj_t * btn){
    child_index++;
    showAddress(child_index, change);
    return LV_RES_OK;
}
static lv_res_t prevCallback(lv_obj_t * btn){
    if(child_index > 0){
        child_index--;
        showAddress(child_index, change);
    }
    return LV_RES_OK;
}
static lv_res_t changeCallback(lv_obj_t * btn){
    change = !change;
    showAddress(child_index, change);
    return LV_RES_OK;
}
static lv_res_t xpubCallback(lv_obj_t * btn){
    titleLbl.text("Your master public key");
    qr.text(xpub.toString());
    dataLbl.text(xpub.toString());
    qr.align(ALIGN_CENTER);
    dataLbl.align(ALIGN_CENTER);
    return LV_RES_OK;
}

void showAddress(unsigned int child_index, bool change){
    stringstream title;
    title << "Your ";
    if(change){
        title << "change ";
    }else{
        title << "receiving ";
    }
    title << "address #" << child_index << ":";
    titleLbl.text(title.str());
    string address = xpub.child(change).child(child_index).address();
    dataLbl.text(address);
    dataLbl.align(ALIGN_CENTER);
    qr.text(string("bitcoin:") + address);
    qr.align(ALIGN_CENTER);
}

void showAddressScreen(){
    gui.clear();
    titleLbl = Label("Your address");
    titleLbl.size(gui.width(), 20);
    titleLbl.position(0, 40);
    titleLbl.alignText(ALIGN_TEXT_CENTER);

    dataLbl = Label(" ");
    dataLbl.size(gui.width()-100, 100); // full width
    dataLbl.position(50, gui.height()-300);
    dataLbl.alignText(ALIGN_TEXT_CENTER);

    qr = QR(" ");
    qr.size(gui.width()-100);
    qr.position(0, 100);
    qr.align(ALIGN_CENTER);

    Button btn(nextCallback, "Next address");
    btn.size(gui.width()/3-20, 80);
    btn.position(gui.width()*2/3 + 10, gui.height()-100);

    Button btn2(prevCallback, "Previous address");
    btn2.size(gui.width()/3-20, 80);
    btn2.position(10, gui.height()-100);

    Button btn3(changeCallback, "Toggle\nchange");
    btn3.size(gui.width()/3-20, 80);
    btn3.position(gui.width()/3 + 10, gui.height()-100);

    Button btn4(xpubCallback, "Show xpub");
    btn4.size(gui.width()/2-20, 80);
    btn4.position(gui.width()/2+10, gui.height()-200);

    Button btn5(toMenuCallback, "Menu");
    btn5.size(gui.width()/2-20, 80);
    btn5.position(10, gui.height()-200);
}

static lv_res_t toMenuCallback(lv_obj_t * btn){
    showMenu();
    return LV_RES_OK;
}

static lv_res_t showAddressesCallback(lv_obj_t * btn){
    showAddressScreen();
    showAddress(child_index, change);
    return LV_RES_OK;
}
static lv_res_t saveXpubCallback(lv_obj_t * btn){
    // TODO: 
    // - check if SD card present
    // - save xpub to the "xpub.txt"
    // - check if write was sucessful
    // - show success / fail message
    if(!SD.detected()){
        showMessage("Error", "SD card is not present");
        return LV_RES_OK;
    }
    int res = SD.save("xpub.txt", xpub.toString());
    if(res != SD_SUCCESS){
        showMessage("Error", "Something wrong with SD card");
        return LV_RES_OK;
    }
    showMessage("Success","Master public key is saved to the SD card");
    return LV_RES_OK;
}
static lv_res_t signConfirmCallback(lv_obj_t * btn){    
    if(!SD.detected()){
        showMessage("Error", "SD card is not present");
        return LV_RES_OK;
    }
    psbt.sign(account);

    uint8_t * raw = (uint8_t *)calloc(psbt.length(), sizeof(uint8_t));
    size_t len = psbt.serialize(raw, psbt.length());
    psbt = PSBT();
    // cout << len << endl;
    char * b64 = (char *)calloc(2*len+1, sizeof(char));
    len = toBase64(raw, len, b64, 2*len+1);
    free(raw);

    int res = SD.save("signed.psbt", b64);
    free(b64);
    if(res != SD_SUCCESS){
        showMessage("Error", "Something wrong with SD card");
        return LV_RES_OK;
    }
    showMessage("Success", "Signed transaction saved to the SD card");
    return LV_RES_OK;
}
void showSignRequest(){
    stringstream ss;
    ss << "Sending:\n\n";
    for(unsigned int i=0; i<psbt.tx.outputsNumber; i++){
        // TODO: better to put it into PSBT method like check_output or something
        bool mine = false;
        if(psbt.txOutsMeta[i].derivationsLen > 0){
            for(unsigned int j=0; j<psbt.txOutsMeta[i].derivationsLen; j++){
                HDPrivateKey prv = account.derive(psbt.txOutsMeta[i].derivations[j].derivation, psbt.txOutsMeta[i].derivations[j].derivationLen);
                PublicKey pub = prv.publicKey();
                // TODO: fix public key comparison
                if(pub.toString() == psbt.txOutsMeta[i].derivations[j].pubkey.toString()){
                    mine = true;
                }
            }
        }
        ss << psbt.tx.txOuts[i].address(&Testnet);
        if(mine){
            ss << " (change)";
        }
        ss << ": " << (float(psbt.tx.txOuts[i].amount)/1e5) << " mBTC\n\n";
    }
    ss << "Fee: " << (float(psbt.fee())/1e5) << " mBTC";

    gui.clear();
    titleLbl = Label("Sign transaction?");
    titleLbl.size(gui.width(), 20);
    titleLbl.position(0, 40);
    titleLbl.alignText(ALIGN_TEXT_CENTER);

    dataLbl = Label(ss.str());
    dataLbl.size(gui.width()-100, 100);
    dataLbl.position(50, 300);
    dataLbl.alignText(ALIGN_TEXT_CENTER);

    Button btn(toMenuCallback, "Cancel");
    btn.size(gui.width()/2-45, 80);
    btn.position(30, gui.height()-100);

    Button btn2(signConfirmCallback, "Confirm");
    btn2.size(gui.width()/2-45, 80);
    btn2.position(gui.width()/2+30, gui.height()-100);
}
static lv_res_t signPSBTCallback(lv_obj_t * btn){
    if(!SD.detected()){
        showMessage("Error", "SD card is not present");
        return LV_RES_OK;
    }
    string s = SD.read("unsigned.psbt");
    if(s.length() == 0){
        showMessage("Fail", "Can't read unsigned.psbt file");
        return LV_RES_OK;
    }
    uint8_t raw[1000];
    unsigned int len = fromBase64(s.c_str(), s.length(), raw, sizeof(raw));
    if(len == 0){
        showMessage("Fail", "Can't convert from base64");
        return LV_RES_OK;
    }
    psbt.reset();
    len = psbt.parse(raw, len);
    if(len == 0){
        showMessage("Fail", "Can't parse PSBT");
        return LV_RES_OK;
    }
    showSignRequest();
    return LV_RES_OK;
}
static lv_res_t wipeCallback(lv_obj_t * btn){
    wipe();
    return LV_RES_OK;
}

/*************** mnemonic stuff ***************/

static lv_res_t newMnemonicCallback(lv_obj_t * btn){
    // TODO: 
    // - generate random buffer (16 or 32 bytes)
    // - create a new mnemonic from it
    // - save this mnemonic
    // - display it to the user
    uint8_t randomBuffer[16];
    getRandomBuffer(randomBuffer, sizeof(randomBuffer));
    string mnemonic = generateMnemonic(randomBuffer, sizeof(randomBuffer));
    saveMnemonic(mnemonic);
    showMessage("Write down your recovery phrase:", mnemonic);
    initKeys(mnemonic);
    return LV_RES_OK;
}

static lv_res_t showMnemonicCallback(lv_obj_t * btn){
    string mnemonic = loadMnemonic();
    showMessage("Here is your recovery phrase:", mnemonic);
    return LV_RES_OK;
}

static lv_res_t enterMnemonicCallback(lv_obj_t * btn){
    return LV_RES_OK;
}

void initKeys(const string mnemonic, const string password){
    // TODO:
    // - derive root key from the mnemonic and empty password
    // - derive account key using m/84'/1'/0'/ derivation path
    // - get account master public key
    root.fromMnemonic(mnemonic, password);
    account = root.derive("m/84'/1'/0'/");
    xpub = account.xpub();
}