From dbc75352ac0566d9dc789a36a4a557f4642dc155 Mon Sep 17 00:00:00 2001 From: kirbylife Date: Wed, 26 Apr 2023 21:50:18 -0600 Subject: [PATCH] Finally works! --- Cargo.lock | 36 ++++---- Cargo.toml | 9 +- src/button.rs | 63 ++++++++++++++ src/crypto.rs | 10 +-- src/datetime.rs | 116 +++++++++++-------------- src/main.rs | 222 +++++++++++++++++++++++++++++++++++++++--------- src/screen.rs | 60 ++++++++++++- src/storage.rs | 131 ++++++++++++++++++++++++++++ tools/menu.py | 157 ++++++++++++++++++++++++++++++++++ upload.sh | 3 +- 10 files changed, 664 insertions(+), 143 deletions(-) create mode 100644 src/button.rs create mode 100644 src/storage.rs create mode 100644 tools/menu.py diff --git a/Cargo.lock b/Cargo.lock index bb9e032..c9ee47a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5,7 +5,7 @@ version = 3 [[package]] name = "arduino-hal" version = "0.1.0" -source = "git+https://github.com/rahix/avr-hal#fb609ab0d14c5a0a44e2dff3e5e514cb612a529a" +source = "git+https://github.com/rahix/avr-hal#e29a7be5b1fa9d490ded11e539d345af041ccef6" dependencies = [ "atmega-hal", "avr-device", @@ -19,7 +19,7 @@ dependencies = [ [[package]] name = "atmega-hal" version = "0.1.0" -source = "git+https://github.com/rahix/avr-hal#fb609ab0d14c5a0a44e2dff3e5e514cb612a529a" +source = "git+https://github.com/rahix/avr-hal#e29a7be5b1fa9d490ded11e539d345af041ccef6" dependencies = [ "avr-device", "avr-hal-generic", @@ -33,9 +33,9 @@ checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" [[package]] name = "avr-device" -version = "0.5.0" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "deb45711b7227cb3f799c40c4e0e8beea1765acff8d06a14bb3479c1a549d8b0" +checksum = "644690f9699830f0efce38ca9d41c7ad51a6ec1829364dd272638938c953673a" dependencies = [ "avr-device-macros", "bare-metal", @@ -45,9 +45,9 @@ dependencies = [ [[package]] name = "avr-device-macros" -version = "0.5.0" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87c428a2384981d8460668c48c766ff18e7e91d2d467f7b103151984ee952d38" +checksum = "ba8b49941a5148cec432d5180d205871f0132f9f15b7225a9f3ad9e2bae026cb" dependencies = [ "proc-macro2", "quote", @@ -57,7 +57,7 @@ dependencies = [ [[package]] name = "avr-hal-generic" version = "0.1.0" -source = "git+https://github.com/rahix/avr-hal#fb609ab0d14c5a0a44e2dff3e5e514cb612a529a" +source = "git+https://github.com/rahix/avr-hal#e29a7be5b1fa9d490ded11e539d345af041ccef6" dependencies = [ "avr-device", "cfg-if 0.1.10", @@ -76,12 +76,6 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8fe8f5a8a398345e52358e18ff07cc17a568fbca5c6f73873d3a62056309603" -[[package]] -name = "binascii" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "383d29d513d8764dcdc42ea295d979eb99c3c9f00607b3692cf68a431f7dca72" - [[package]] name = "cfg-if" version = "0.1.10" @@ -107,8 +101,6 @@ dependencies = [ [[package]] name = "ds323x" version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "057e63d03a3beb83d3c9b7a19b799c60e203b0c53f1bbec6a715804907f73a2b" dependencies = [ "embedded-hal", "rtcc", @@ -184,9 +176,9 @@ checksum = "dc375e1527247fe1a97d8b7156678dfe7c1af2fc075c9a4db3690ecd2a148068" [[package]] name = "proc-macro2" -version = "1.0.54" +version = "1.0.56" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e472a104799c74b514a57226160104aa483546de37e839ec50e3c2e41dd87534" +checksum = "2b63bdb0cd06f1f4dedf69b254734f9b45af66e4a031e42a7480257d9898b435" dependencies = [ "unicode-ident", ] @@ -200,6 +192,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "randomize" +version = "3.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88c38c99b51f33c9fcc655252bf02ac8048eb70f35244e4697b0de9c473e940a" + [[package]] name = "rtcc" version = "0.3.0" @@ -220,9 +218,11 @@ name = "rustytoken" version = "0.1.0" dependencies = [ "arduino-hal", - "binascii", "ds323x", + "embedded-hal", + "nb 1.1.0", "panic-halt", + "randomize", "sha1_smol", "ufmt 0.2.0", ] diff --git a/Cargo.toml b/Cargo.toml index 96518d8..14336d6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,14 +7,15 @@ edition = "2021" arduino-hal = { git = "https://github.com/rahix/avr-hal", features = ["arduino-uno"] } panic-halt = "0.2.0" sha1_smol = "1.0.0" -binascii = { version = "0.1", default-features = false, features = ["decode"] } -ds323x = "0.5.0" -ufmt = "0.2.0" +ds323x = { path = "../ds323x-rs" } embedded-hal = "0.2.7" +nb = "1.1.0" +randomize = "3.0.1" +ufmt = "0.2.0" [profile.release] lto = true panic = "abort" strip = true -opt-level = "z" +opt-level = "s" codegen-units = 1 diff --git a/src/button.rs b/src/button.rs new file mode 100644 index 0000000..6dce4e6 --- /dev/null +++ b/src/button.rs @@ -0,0 +1,63 @@ +use embedded_hal::digital::v2::InputPin; + +pub struct Button<'a, P> { + pin: &'a P, + state: bool, + is_pullup: bool, +} + +#[derive(PartialEq)] +pub enum Event { + PressDown, + PressUp, + Pressed, + None, +} + +impl<'a, P: InputPin> Button<'a, P> { + pub fn new(pin: &'a P, is_pullup: bool) -> Self { + let mut state = pin.is_high().map_err(|_| ()).unwrap(); + if is_pullup { + state = !state; + } + + Button { + pin, + state, + is_pullup, + } + } + + fn read_pin(&self) -> bool { + let mut inc = 0; + + for _ in 0..100 { + let mut state = self.pin.is_high().map_err(|_| ()).unwrap(); + if self.is_pullup { + state = !state; + } + if state { + inc += 1; + } + } + + inc > 70 + } + + pub fn update(&mut self) -> Event { + let actual_state = self.read_pin(); + + match (self.state, actual_state) { + (false, false) => Event::None, + (true, false) => { + self.state = !self.state; + Event::PressUp + } + (false, true) => { + self.state = !self.state; + Event::PressDown + } + (true, true) => Event::Pressed, + } + } +} diff --git a/src/crypto.rs b/src/crypto.rs index 77c3377..259b648 100644 --- a/src/crypto.rs +++ b/src/crypto.rs @@ -5,14 +5,8 @@ const BLOCK_SIZE: usize = 64; pub fn hmac_sha1(key: &[u8], message: &[u8], output: &mut [u8]) { // Preprocess the key. let mut key_padded = [0u8; BLOCK_SIZE]; - if key.len() > BLOCK_SIZE { - let mut hash = Sha1::new(); - hash.update(key); - let key_digest = hash.digest().bytes(); - key_padded[..20].copy_from_slice(&key_digest); - } else { - key_padded[..key.len()].copy_from_slice(key); - } + key_padded[..key.len()].copy_from_slice(key); + for byte in &mut key_padded { *byte ^= 0x36; } diff --git a/src/datetime.rs b/src/datetime.rs index 0f9539d..2b181c7 100644 --- a/src/datetime.rs +++ b/src/datetime.rs @@ -13,15 +13,7 @@ impl Datetime { pub fn unix_epoch(&self) -> u64 { let mut year: i64 = self.year as i64 - 1900; let mut month: i64 = self.month as i64 - 1; - if month >= 12 || month < 0 { - let mut adj = month / 12; - month %= 12; - if month < 0 { - adj -= 1; - month += 12; - } - year += adj as i64; - } + let (mut t, is_leap) = year_to_secs(year); t += month_to_secs(month, is_leap); t += 86400 * (self.day as i64 - 1); @@ -35,21 +27,21 @@ impl Datetime { let year = rtc.year().map_err(|_| ()).unwrap(); let month = rtc.month().map_err(|_| ()).unwrap(); let day = rtc.day().map_err(|_| ()).unwrap(); - let hour = match rtc.hours().map_err(|_| ()).unwrap() { + let hours = match rtc.hours().map_err(|_| ()).unwrap() { ds323x::Hours::AM(val) => val, ds323x::Hours::PM(val) => val + 12, ds323x::Hours::H24(val) => val, }; - let minute = rtc.minutes().map_err(|_| ()).unwrap(); - let second = rtc.seconds().map_err(|_| ()).unwrap(); + let minutes = rtc.minutes().map_err(|_| ()).unwrap(); + let seconds = rtc.seconds().map_err(|_| ()).unwrap(); Datetime { - year: year, - month: month, - day: day, - hours: hour, - minutes: minute, - seconds: second, + year, + month, + day, + hours, + minutes, + seconds, } } } @@ -58,64 +50,52 @@ fn year_to_secs(year: i64) -> (i64, bool) { let is_leap: bool; let res: i64; - if year - 2 <= 136 { - let y = year as i32; - let mut leaps = (y - 68) >> 2; - if !((y - 68) & 3 != 0) { - leaps -= 1; - is_leap = true; - } else { - is_leap = false; - } - res = 31536000 * (y - 70) as i64 + 86400 * leaps as i64; - } else { - let mut cycles: i64; - let centuries: i64; - let mut leaps: i64; - let mut rem: i64; + let mut cycles: i64; + let centuries: i64; + let mut leaps: i64; + let mut rem: i64; - cycles = (year - 100) / 400; - rem = (year - 100) % 400; - if rem < 0 { - cycles -= 1; - rem += 400; + cycles = (year - 100) / 400; + rem = (year - 100) % 400; + if rem < 0 { + cycles -= 1; + rem += 400; + } + if rem == 0 { + is_leap = true; + centuries = 0; + leaps = 0; + } else { + if rem >= 200 { + if rem >= 300 { + centuries = 3; + rem -= 300; + } else { + centuries = 2; + rem -= 200; + } + } else { + if rem >= 100 { + centuries = 1; + rem -= 100; + } else { + centuries = 0; + } } if rem == 0 { - is_leap = true; - centuries = 0; + is_leap = false; leaps = 0; } else { - if rem >= 200 { - if rem >= 300 { - centuries = 3; - rem -= 300; - } else { - centuries = 2; - rem -= 200; - } - } else { - if rem >= 100 { - centuries = 1; - rem -= 100; - } else { - centuries = 0; - } - } - if rem == 0 { - is_leap = false; - leaps = 0; - } else { - leaps = rem / 4; - rem %= 4; - is_leap = rem == 0; - } + leaps = rem / 4; + rem %= 4; + is_leap = rem == 0; } - - leaps += 97 * cycles + 24 * centuries - (is_leap as i64); - - res = (year - 100) * 31536000 + leaps * 86400 + 946684800 + 86400; } + leaps += 97 * cycles + 24 * centuries - (is_leap as i64); + + res = (year - 100) * 31536000 + leaps * 86400 + 946684800 + 86400; + (res, is_leap) } diff --git a/src/main.rs b/src/main.rs index bebe2e6..5e330c6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,53 +1,46 @@ #![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::hmac_sha1; - +use ds323x::Rtcc; +use nb::block; use panic_halt as _; +use storage::SECRET_KEY_MAX_LEN; +use storage::SECRET_KEY_NAME_LEN; const INTERVAL: u64 = 30; -const SECRET_KEY_MAX_LEN: usize = 32; -fn process_secret<'a>(secret: &str, output: &mut [u8]) -> usize { - // RFC 6238 Doesn't specify a max length for the secret keys, So I chose 32 - // You can change the max length by modifying the const - let mut result = [0u8; SECRET_KEY_MAX_LEN]; - let mut inc = 0; - let offset = match 8 - (secret.len() % 8) { - 8 => 0, - n => n, - }; +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 WIPE_TOKENS: u8 = 50; +const EXIT: u8 = 254; - for c in secret.chars() { - result[inc] = c.to_ascii_uppercase() as u8; - inc += 1; - } - for _ in 0..offset { - result[inc] = '=' as u8; - inc += 1; - } - - let decoded = match binascii::b32decode(&result[..inc], output).ok() { - Some(bytes) => bytes, - None => "".as_bytes(), // TODO: Implement a way to display errors - }; - decoded.len() +fn decimal_to_packed_bcd(dec: u8) -> u8 { + ((dec / 10) << 4) | (dec % 10) } -fn generate_otp_token(private_key: &str, actual_time: u64) -> u32 { - let mut output = [0u8; 32]; - let private_key_len = process_secret(private_key, &mut output); - +fn generate_otp_token(private_key: &[u8], actual_time: u64) -> u32 { let interval = actual_time / INTERVAL; let msg = interval.to_be_bytes(); let mut hmac_digest = [0u8; 20]; - hmac_sha1(&output[..private_key_len], &msg, &mut hmac_digest); + hmac_sha1(private_key, &msg, &mut hmac_digest); let start = (hmac_digest[19] & 0xF) as usize; let bytes: [u8; 4] = hmac_digest[start..start + 4].try_into().unwrap(); @@ -72,21 +65,170 @@ fn main() -> ! { ); let mut rtc = ds323x::Ds323x::new_ds3231(i2c_clock); + let mut rs = pins.d12.into_output(); + let mut en = pins.d11.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 datetime = datetime::Datetime::from_ds3231(&mut rtc); - let timestamp = datetime.unix_epoch(); - ufmt::uwriteln!(&mut serial, "Timestamp: {}", timestamp).unwrap(); + let mut eeprom = arduino_hal::Eeprom::new(dp.EEPROM); + let mut tokens = storage::Tokens::new(&mut eeprom, datetime.unix_epoch()); - let token = generate_otp_token("FS7J22EHLLSOGKUVJKV2XE7FTIX24JAJ", 1680155315); - ufmt::uwriteln!(&mut serial, "token: {}", token).unwrap(); + let up = pins.d6.into_pull_up_input(); + let mut button = button::Button::new(&up, true); - delay_ms(5000); - let token = generate_otp_token("FS7J22EHLLSOGKUVJKV2XE7FTIX24JAJ", 1680155315); - ufmt::uwriteln!(&mut serial, "token: {}", token).unwrap(); + // button.update(); + // if button.update() == button::Event::Pressed || true { + // display.write_str("Connected to"); + // display.set_cursor(0, 1); + // display.write_str("USB..."); + // loop { + // // Waiting from a command since the tool + // let cmd = block!(serial.read()).unwrap_or(u8::MAX); + // display.set_cursor(8, 1); + // display.write_str("---"); + // display.set_cursor(8, 1); + // display.write_u32(cmd as u32); - delay_ms(5000); - let token = generate_otp_token("FS7J22EHLLSOGKUVJKV2XE7FTIX24JAJ", 1680155315); - ufmt::uwriteln!(&mut serial, "token: {}", token).unwrap(); + // match cmd { + // EXIT => {} + // HANDSHAKE => serial.write(OK).unwrap(), + // SET_TIMESTAMP => { + // serial.write(OK).unwrap(); + // let year = loop { + // let y = block!(serial.read()).unwrap(); + // serial.write(y).unwrap(); + // if block!(serial.read()).unwrap() == OK { + // break y; + // } + // }; + // let month = loop { + // let m = block!(serial.read()).unwrap(); + // serial.write(m).unwrap(); + // if block!(serial.read()).unwrap() == OK { + // break m; + // } + // }; + // let day = loop { + // let d = block!(serial.read()).unwrap(); + // serial.write(d).unwrap(); + // if block!(serial.read()).unwrap() == OK { + // break d; + // } + // }; + // let hours = loop { + // let h = block!(serial.read()).unwrap(); + // serial.write(h).unwrap(); + // if block!(serial.read()).unwrap() == OK { + // break h; + // } + // }; + // let minutes = loop { + // let m = block!(serial.read()).unwrap(); + // serial.write(m).unwrap(); + // if block!(serial.read()).unwrap() == OK { + // break m; + // } + // }; + // let seconds = loop { + // let s = block!(serial.read()).unwrap(); + // serial.write(s).unwrap(); + // if block!(serial.read()).unwrap() == OK { + // break s; + // } + // }; + // 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) { + // Some(_) => serial.write(OK).unwrap(), + // None => serial.write(ERROR).unwrap(), + // } + // } + // None => serial.write(ERROR).unwrap(), + // }, + // WIPE_TOKENS => { + // serial.write(OK).unwrap(); + // tokens.wipe_all_tokens(); + // serial.write(OK).unwrap(); + // } + // _ => {} + // u8::MAX => {} + // } + // } + // } + // loop {} + + 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) { + Some((len_name, len_key)) => { + let name = core::str::from_utf8(&buff_name[..len_name]).unwrap(); + let key = &buff_key[..len_key]; + let dt = datetime::Datetime::from_ds3231(&mut rtc); + + let datetime::Datetime { + year, + month, + day, + hours, + minutes, + seconds, + } = dt; + + let timestamp = dt.unix_epoch(); + + let token = generate_otp_token(key, timestamp); + + display.set_cursor(0, 0); + display.write_str(name); + display.set_cursor(0, 1); + display.write_u32(token); + } + _ => { + // display.clear(); + // display.set_cursor(0, 0); + // display.write_str("ERROR"); + } + } + } + _ => {} + } loop { led.toggle(); diff --git a/src/screen.rs b/src/screen.rs index 6062191..cc64ad5 100644 --- a/src/screen.rs +++ b/src/screen.rs @@ -2,11 +2,23 @@ use arduino_hal::delay_ms; use arduino_hal::delay_us; use embedded_hal::digital::v2::OutputPin; +// The CMDs codes are documented here: +// https://www.openhacks.com/uploadsproductos/eone-1602a1.pdf const INIT_CMD: u8 = 0x03; // 0b00000011 const MODE_4_BITS_CMD: u8 = 0x02; // 0b00000010 -const DISPLAY_MODE_CMD: u8 = 0x28; // 0b00101000 +const DISPLAY_MODE_16X2_CMD: u8 = 0x28; // 0b00101000 const ENABLE_SCREEN_CMD: u8 = 0x0C; // 0b00001100 const CLEAN_SCREEN_CMD: u8 = 0x01; // 0b00000001 +const SHIFT_LEFT_CMD: u8 = 0x10; // 0b10000 +const SHIFT_RIGHT_CMD: u8 = 0x14; // 0b10100 +const MOVE_CURSOR_CMD: u8 = 0x80; // 0b10000000 + +const ROW_OFFSETS: [u8; 2] = [0x00, 0x40]; + +pub enum Direction { + Left = SHIFT_LEFT_CMD as isize, + Right = SHIFT_RIGHT_CMD as isize, +} pub struct StnScreen<'a, RS, EN, D4, D5, D6, D7> { rs: &'a mut RS, @@ -55,8 +67,9 @@ impl< delay_us(100); display.write_nibble(MODE_4_BITS_CMD, false); delay_us(100); - display.send_cmd(DISPLAY_MODE_CMD, false); + display.send_cmd(DISPLAY_MODE_16X2_CMD, false); display.send_cmd(ENABLE_SCREEN_CMD, false); + display.clear(); display } @@ -85,8 +98,19 @@ impl< self.pulse_enable(); } - fn clear(&mut self) { - display.send_cmd(CLEAN_SCREEN_CMD, false); + pub fn clear(&mut self) { + self.send_cmd(CLEAN_SCREEN_CMD, false); + } + + pub fn shift(&mut self, dir: Direction) { + let cmd = dir as u8; + self.send_cmd(cmd, false); + } + + pub fn set_cursor(&mut self, col: u8, row: u8) { + let row_offset = ROW_OFFSETS[row as usize]; + let address_position = row_offset | col; + self.send_cmd(MOVE_CURSOR_CMD | address_position, false); } pub fn write_char(&mut self, c: char) { @@ -102,8 +126,36 @@ impl< self.write_char(c); } } + + pub fn write_u32(&mut self, num: u32) { + // The max number in an u32 is: 4,294,967,295 + let mut buff = [0; 10]; + let len = num_chars(num, &mut buff); + self.write_str(&core::str::from_utf8(&buff[..len]).unwrap()); + } } fn last_bit(n: u8) -> bool { (n & 1) != 0 } + +fn num_rev(mut num: u32) -> u32 { + let mut rev_num = 0; + while num > 0 { + let digit = num % 10; + rev_num = rev_num * 10 + digit; + num /= 10; + } + rev_num +} + +fn num_chars(n: u32, output: &mut [u8]) -> usize { + let mut n = num_rev(n); + let mut i = 0; + while n > 0 { + output[i] = (n % 10) as u8 + '0' as u8; + n = n / 10; + i += 1; + } + i +} diff --git a/src/storage.rs b/src/storage.rs new file mode 100644 index 0000000..a3d88d6 --- /dev/null +++ b/src/storage.rs @@ -0,0 +1,131 @@ +use arduino_hal::eeprom::Eeprom; +use randomize::PCG32; + +pub const SECRET_KEY_MAX_LEN: u16 = 32; +pub const SECRET_KEY_NAME_LEN: u16 = 16; +pub const SECRET_KEY_FULL_LEN: u16 = SECRET_KEY_MAX_LEN + SECRET_KEY_NAME_LEN; + +pub struct Tokens<'a> { + mem: &'a mut Eeprom, + pub current: Option, + capacity: u16, + rand: PCG32, +} + +impl<'a> Tokens<'a> { + pub fn new(mem: &'a mut Eeprom, rand_seed: u64) -> Self { + let capacity = mem.capacity() / SECRET_KEY_FULL_LEN; + + let mut tokens = Tokens { + mem, + capacity, + rand: PCG32::seed(rand_seed, 1), + current: None, + }; + tokens.current = tokens.first(); + + tokens + } + + pub fn search_free(&self) -> Option { + for n in 0..self.capacity { + let index = SECRET_KEY_FULL_LEN * n; + if self.mem.read_byte(index) == 255 { + return Some(n); + } + } + + None + } + + fn first(&self) -> Option { + for n in 0..self.capacity { + let index = SECRET_KEY_FULL_LEN * n; + if self.mem.read_byte(index) != 255 { + return Some(n); + } + } + return None; + } + + pub fn next(&mut self) -> Option { + let mut index = self.current.unwrap(); + for _ in 0..self.capacity { + index = (index + 1) % self.capacity; + if self.mem.read_byte(index * SECRET_KEY_FULL_LEN) != 255 { + self.current = Some(index); + return Some(index); + } + } + None + } + + pub fn prev(&mut self) -> Option { + todo!(); + } + + pub fn read(&self, index: u16, name: &mut [u8], key: &mut [u8]) -> Option<(usize, usize)> { + name.fill(0); + key.fill(0); + let index_name = index * SECRET_KEY_FULL_LEN; + let index_key = (index * SECRET_KEY_FULL_LEN) + SECRET_KEY_NAME_LEN; + match self.mem.read(index_name, name) { + Ok(n) => {} + Err(_) => return None, + } + match self.mem.read(index_key, key) { + Ok(n) => {} + Err(_) => return None, + } + let mut len_name = name + .iter() + .position(|&n| n == 0) + .unwrap_or(SECRET_KEY_NAME_LEN as usize); + let mut len_key = key + .iter() + .position(|&n| n == 0) + .unwrap_or(SECRET_KEY_MAX_LEN as usize); + Some((len_name, len_key)) + } + + pub fn write(&mut self, index: u16, name: &[u8], key: &[u8]) -> Option { + let index_name = index * SECRET_KEY_FULL_LEN; + let index_key = (index * SECRET_KEY_FULL_LEN) + SECRET_KEY_NAME_LEN; + match self.mem.write(index_name, name) { + Ok(_) => {} + Err(_) => return None, + } + match self.mem.write(index_key, key) { + Ok(_) => {} + Err(_) => return None, + } + Some(index) + } + + pub fn delete(&mut self, index: u16) -> Option { + // The Arduino's EEPROM memory has a maximum number of write cycles. + // To keep writes to a minimum, only the first byte of the token name is set to 0 + // and a byte of the key is randomly chosen to be overwritten with + // another random value, so that it's unrecoverable. + let index_name = index * SECRET_KEY_FULL_LEN; + let index_key = (index * SECRET_KEY_FULL_LEN) + SECRET_KEY_NAME_LEN; + let index_key = index_key + (self.rand.next_u32() % SECRET_KEY_MAX_LEN as u32) as u16; + let rand_byte = (self.rand.next_u32() % 255) as u8; + self.mem.write_byte(index_name, 255); + self.mem.write_byte(index_key, rand_byte); + Some(index) + } + + pub fn wipe_all_tokens(&mut self) -> u8 { + let mut inc = 0; + + for n in 0..self.capacity { + let index = SECRET_KEY_FULL_LEN * n; + if self.mem.read_byte(index) != 255 { + self.delete(index); + inc += 1; + } + } + inc + } +} diff --git a/tools/menu.py b/tools/menu.py new file mode 100644 index 0000000..71510e4 --- /dev/null +++ b/tools/menu.py @@ -0,0 +1,157 @@ +from datetime import datetime +from time import sleep +from typing import Any +import base64 +import serial +import sys + +ENDL = bytes([0]) +OK = bytes([1]) +ERROR = bytes([2]) +HANDSHAKE = bytes([3]) +SET_TIMESTAMP = bytes([10]) +ADD_TOKEN = bytes([20]) +DELETE_TOKEN = bytes([30]) +GET_TOKENS = bytes([40]) +WIPE_TOKENS = bytes([50]) +EXIT = bytes([254]) + +def loop_input(msg: str, valid_values: Any) -> str: + valid_values = list(map(str, valid_values)) + + while True: + data = input(msg) + + if data not in valid_values: + print(f"'{data}' isn't a valid value, please enter a value in {valid_values}") + else: + break + return data + +def b(n: int) -> bytes: + return bytes([n]) + +def process_secret(secret: str) -> bytes: + offset = 8 - (len(secret) % 8) + if offset != 8: + secret += "=" * offset + return base64.b32decode(secret, casefold=True) + +def send_and_validate(value: bytes, conn): + while True: + print(f"Send {list(value)}") + conn.write(value) + resp = conn.read() + print(f"Received {list(resp)}") + if resp == value: + conn.write(OK) + break + else: + conn.write(ERROR) + +def get_datetime_items() -> dict[str, bytes]: + now = datetime.utcnow() + + return { + "year": b(now.year - 2000), + "month": b(now.month), + "day": b(now.day), + "hours": b(now.hour), + "minutes": b(now.minute), + "seconds": b(now.second) + } + +def main(argv: list[str]): + port = argv[-1] + + conn = serial.Serial(port=port, baudrate=9600) + print("UP + Reset to enable the USB connection") + input("Press Enter if the device in the USB mode...") + + conn.write(HANDSHAKE) + sleep(0.1) + res = conn.read() + if res != OK: + print("A handshake could not be performed") + print("Check the connection with your Arduino") + return 1 + else: + print("Handshake successfully performed") + + while True: + print("What do you want to do?") + print("1) Update timestamp") + print("2) Add a new token") + print("3) Remove a token") + print("4) WIPE ALL THE TOKENS") + print("5) EXIT") + + opt = loop_input(">>> ", range(1, 6)) + + # Update Timestamp + if opt == "1": + conn.write(SET_TIMESTAMP) + sleep(0.1) + resp = conn.read() + sleep(0.1) + if resp == ERROR: + print(f"Error in the communication: Error {resp}") + continue + date_items = get_datetime_items(); + send_and_validate(date_items["year"], conn) + send_and_validate(date_items["month"], conn) + send_and_validate(date_items["day"], conn) + send_and_validate(date_items["hours"], conn) + send_and_validate(date_items["minutes"], conn) + send_and_validate(date_items["seconds"], conn) + resp = conn.read() + if resp != OK: + print(f"Error in the communication: Error {resp}") + else: + print("Timestamp updated successfully!") + # Add token + elif opt == "2": + conn.write(ADD_TOKEN) + sleep(0.1) + resp = conn.read() + if resp == ERROR: + print("The memory of the device is full") + continue + name = input("Enter the name of the new token (16 chars max):\n>>> ") + name = name.strip()[:16].encode("ascii") + + key = input("Enter the OTP secret key (32 chars max):\n>>> ") + key = process_secret(key.strip()[:32]) + + for ch in name: + conn.write(b(ch)) + conn.write(ENDL) + for ch in key: + conn.write(b(ch)) + conn.write(ENDL) + + resp = conn.read() + if resp == ERROR: + print("Error trying to add the token, try again") + else: + print("Token added successfully!") + # Wipe tokens + elif opt == "4": + conn.write(WIPE_TOKENS) + sleep(0.1) + _ = conn.read() + resp = conn.read() + if resp == OK: + print("All the tokens wipped successfully!") + elif opt == "5": + return 0 + +if __name__ == "__main__": + attrs = sys.argv + if len(attrs) == 1: + print("You need to specify the Arduino's serial port") + print("Example:") + print(f"python {attrs[0]} /dev/ttyUSB0") + sys.exit(1) + + sys.exit(main(attrs[1:])) diff --git a/upload.sh b/upload.sh index 65248a1..96d413b 100755 --- a/upload.sh +++ b/upload.sh @@ -1 +1,2 @@ -avrdude -p atmega328p -c arduino -P /dev/ttyACM0 -U flash:w:target/avr-atmega328p/release/rustytoken.elf:e +# avrdude -p atmega328p -c arduino -P /dev/ttyACM0 -U flash:w:target/avr-atmega328p/release/rustytoken.elf:e +avrdude -p atmega328p -c arduino -P /dev/ttyUSB0 -U flash:w:target/avr-atmega328p/release/rustytoken.elf:e