rustytoken/src/main.rs

#![no_std]
#![no_main]

mod button;
mod crypto;
mod datetime;
mod screen;
mod storage;

use arduino_hal::default_serial;
use arduino_hal::delay_ms;
// use arduino_hal::delay_ms;
use arduino_hal::prelude::*;
use crypto::generate_otp_token;
use crypto::INTERVAL;
use ds323x::Rtcc;
use nb::block;
use panic_halt as _;
use storage::SECRET_KEY_MAX_LEN;
use storage::SECRET_KEY_NAME_LEN;

const ENDL: u8 = 0;
const OK: u8 = 1;
const ERROR: u8 = 2;
const HANDSHAKE: u8 = 3;
const SET_TIMESTAMP: u8 = 10;
const ADD_TOKEN: u8 = 20;
const DELETE_TOKEN: u8 = 30;
const GET_TOKENS: u8 = 40;
const SOFT_WIPE_TOKENS: u8 = 50;
const HARD_WIPE_TOKENS: u8 = 60;
const EXIT: u8 = 255;

#[arduino_hal::entry]
fn main() -> ! {
    let dp = arduino_hal::Peripherals::take().unwrap();
    let pins = arduino_hal::pins!(dp);

    let mut led = pins.d13.into_output();
    let mut serial = default_serial!(dp, pins, 9600);

    let mut rs = pins.d9.into_output();
    let mut en = pins.d8.into_output();
    let mut d4 = pins.d5.into_output();
    let mut d5 = pins.d4.into_output();
    let mut d6 = pins.d3.into_output();
    let mut d7 = pins.d2.into_output();

    let mut display = screen::StnScreen::new(&mut rs, &mut en, &mut d4, &mut d5, &mut d6, &mut d7);
    display.clear();

    let i2c_clock = arduino_hal::I2c::new(
        dp.TWI,
        pins.a4.into_pull_up_input(), // SDA
        pins.a5.into_pull_up_input(), // SCL
        50000,
    );
    let mut rtc = ds323x::Ds323x::new_ds3231(i2c_clock);

    let mut eeprom = arduino_hal::Eeprom::new(dp.EEPROM);
    let mut tokens = storage::Tokens::new(&mut eeprom);

    let up = pins.a2.into_pull_up_input();
    let mut up_button = button::Button::new(&up, true);
    let down = pins.a1.into_pull_up_input();
    let mut down_button = button::Button::new(&down, true);

    up_button.update();
    if up_button.update() == button::Event::Pressed || tokens.current.is_none() {
        display.write_str("Connected to");
        display.set_cursor(0, 1);
        display.write_str("USB...");
        loop {
            // Waiting from a command from the tool
            let cmd = block!(serial.read()).unwrap_or(u8::MAX);

            match cmd {
                HANDSHAKE => serial.write(OK).unwrap(),
                SET_TIMESTAMP => {
                    serial.write(OK).unwrap();
                    let year = block!(serial.read()).unwrap();
                    let month = block!(serial.read()).unwrap();
                    let day = block!(serial.read()).unwrap();
                    let hours = block!(serial.read()).unwrap();
                    let minutes = block!(serial.read()).unwrap();
                    let seconds = block!(serial.read()).unwrap();

                    rtc.set_year(year as u16 + 2000).unwrap();
                    rtc.set_month(month).unwrap();
                    rtc.set_day(day).unwrap();
                    rtc.set_hours(ds323x::Hours::H24(hours)).unwrap();
                    rtc.set_minutes(minutes).unwrap();
                    rtc.set_seconds(seconds).unwrap();
                    serial.write(OK).unwrap();
                }
                ADD_TOKEN => match tokens.search_free() {
                    Some(index) => {
                        serial.write(OK).unwrap();
                        let mut name_buff = [0u8; SECRET_KEY_NAME_LEN as usize];
                        for i in 0..name_buff.len() {
                            let ch = block!(serial.read()).unwrap();
                            if ch == ENDL {
                                break;
                            } else {
                                name_buff[i] = ch;
                            }
                        }

                        let mut key_buff = [0u8; SECRET_KEY_MAX_LEN as usize];
                        for i in 0..key_buff.len() {
                            let ch = block!(serial.read()).unwrap();
                            if ch == ENDL {
                                break;
                            } else {
                                key_buff[i] = ch;
                            }
                        }
                        match tokens.write(index, &name_buff, &key_buff) {
                            Ok(_) => serial.write(OK).unwrap(),
                            Err(_) => serial.write(ERROR).unwrap(),
                        }
                    }
                    None => serial.write(ERROR).unwrap(),
                },
                SOFT_WIPE_TOKENS => {
                    serial.write(OK).unwrap();
                    let deleted_tokens = tokens.soft_wipe_all_tokens();
                    serial.write(deleted_tokens).unwrap();
                    serial.write(OK).unwrap();
                }
                HARD_WIPE_TOKENS => {
                    serial.write(OK).unwrap();
                    tokens.hard_wipe_all_tokens();
                    serial.write(OK).unwrap();
                }
                EXIT => break,
                _ => {}
            }
        }
    }

    let mut last_time = 0;
    let mut changed = false;

    loop {
        let timestamp = datetime::Datetime::from_ds3231(&mut rtc).unix_epoch();

        if up_button.update() == button::Event::PressUp {
            changed = true;
            tokens.next();
        }
        if down_button.update() == button::Event::PressUp {
            changed = true;
            tokens.prev();
        }
        let curr_time = INTERVAL - (timestamp % INTERVAL);
        if curr_time == last_time {
            delay_ms(10);
            continue;
        }
        last_time = curr_time;
        display.set_cursor(12, 1);
        display.write_str("  ");
        display.set_cursor(12, 1);
        display.write_u8(curr_time as u8);

        if changed {
            display.set_cursor(0, 0);
            display.write_str("                ");
            changed = false;
        }

        match tokens.current {
            Some(index) => {
                let mut buff_name = [0u8; SECRET_KEY_NAME_LEN as usize];
                let mut buff_key = [0u8; SECRET_KEY_MAX_LEN as usize];
                match tokens.read(index, &mut buff_name, &mut buff_key) {
                    Ok((len_name, len_key)) => {
                        let name = core::str::from_utf8(&buff_name[..len_name]).unwrap();
                        let key = &buff_key[..len_key];

                        let token = generate_otp_token(key, timestamp);

                        display.set_cursor(0, 0);
                        display.write_str(name);
                        display.set_cursor(0, 1);
                        display.write_token(token);
                    }
                    _ => {
                        display.clear();
                        display.set_cursor(0, 0);
                        display.write_str("ERROR");
                    }
                }
            }
            _ => {}
        }
        changed = false;
    }
}