From f83638bfae7b198163535baf3e165bd35618452c Mon Sep 17 00:00:00 2001 From: kirbylife Date: Thu, 27 Apr 2023 15:21:17 -0600 Subject: [PATCH] Implement the menu --- src/main.rs | 167 ++++++++++++++++++++++++------------------------- src/screen.rs | 33 ++++++++++ src/storage.rs | 38 +++++------ tools/menu.py | 24 ++----- 4 files changed, 137 insertions(+), 125 deletions(-) diff --git a/src/main.rs b/src/main.rs index 2934843..8b59bf5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -27,8 +27,9 @@ 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; +const SOFT_WIPE_TOKENS: u8 = 50; +const HARD_WIPE_TOKENS: u8 = 60; +const EXIT: u8 = 255; #[arduino_hal::entry] fn main() -> ! { @@ -62,68 +63,31 @@ fn main() -> ! { let mut tokens = storage::Tokens::new(&mut eeprom, datetime.unix_epoch()); let up = pins.d6.into_pull_up_input(); - let mut button = button::Button::new(&up, true); + let mut up_button = button::Button::new(&up, true); + let down = pins.d7.into_pull_up_input(); + let mut down_button = button::Button::new(&down, true); - button.update(); - if button.update() == button::Event::Pressed { + up_button.update(); + if up_button.update() == button::Event::Pressed { display.write_str("Connected to"); display.set_cursor(0, 1); display.write_str("USB..."); loop { - // Waiting from a command since the tool + // Waiting from a command from 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); 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; - } - }; + 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(); @@ -155,53 +119,84 @@ fn main() -> ! { } } match tokens.write(index, &name_buff, &key_buff) { - Some(_) => serial.write(OK).unwrap(), - None => serial.write(ERROR).unwrap(), + Ok(_) => serial.write(OK).unwrap(), + Err(_) => serial.write(ERROR).unwrap(), } } None => serial.write(ERROR).unwrap(), }, - WIPE_TOKENS => { + SOFT_WIPE_TOKENS => { serial.write(OK).unwrap(); - tokens.wipe_all_tokens(); + tokens.soft_wipe_all_tokens(); serial.write(OK).unwrap(); } + HARD_WIPE_TOKENS => { + serial.write(OK).unwrap(); + tokens.hard_wipe_all_tokens(); + serial.write(OK).unwrap(); + } + EXIT => break, _ => {} } } } - 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 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"); - } - } - } - _ => {} - } + let mut changed = false; + let mut last_index = 100; + let mut last_time = 0; loop { - led.toggle(); - delay_ms(1000); + 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 { + todo!("To be implemented"); + } + 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; } } diff --git a/src/screen.rs b/src/screen.rs index cc64ad5..802fc63 100644 --- a/src/screen.rs +++ b/src/screen.rs @@ -127,7 +127,40 @@ impl< } } + pub fn write_token(&mut self, mut token: u32) { + let mut buff = [0u8; 7]; + for (n, chr) in buff.iter_mut().enumerate() { + *chr = if n == 3 { + '-' as u8 + } else { + let ch = (token % 10) as u8 + '0' as u8; + token /= 10; + ch + } + } + buff.reverse(); + let s = core::str::from_utf8(&buff).unwrap(); + self.write_str(s); + } + + pub fn write_u8(&mut self, mut num: u8) { + if num == 0 { + self.write_char('0'); + return; + } + + let mut buff = [0u8; 3]; + let num_len = num.ilog10() as usize; + for i in (0..=num_len).rev() { + buff[i] = num % 10 + '0' as u8; + num /= 10; + } + self.write_str(core::str::from_utf8(&buff[..=num_len]).unwrap()); + } + pub fn write_u32(&mut self, num: u32) { + // This function is for debugging purposes only, it does not show + // The 0's to the right of the number // The max number in an u32 is: 4,294,967,295 let mut buff = [0; 10]; let len = num_chars(num, &mut buff); diff --git a/src/storage.rs b/src/storage.rs index 7494c65..57f5f0a 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -64,19 +64,15 @@ impl<'a> Tokens<'a> { todo!(); } - pub fn read(&self, index: u16, name: &mut [u8], key: &mut [u8]) -> Option<(usize, usize)> { + pub fn read(&self, index: u16, name: &mut [u8], key: &mut [u8]) -> Result<(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(_) => {} - Err(_) => return None, - } - match self.mem.read(index_key, key) { - Ok(_) => {} - Err(_) => return None, - } + + self.mem.read(index_name, name).map_err(|_| ())?; + self.mem.read(index_key, key).map_err(|_| ())?; + let len_name = name .iter() .position(|&n| n == 0) @@ -85,21 +81,17 @@ impl<'a> Tokens<'a> { .iter() .position(|&n| n == 0) .unwrap_or(SECRET_KEY_MAX_LEN as usize); - Some((len_name, len_key)) + Ok((len_name, len_key)) } - pub fn write(&mut self, index: u16, name: &[u8], key: &[u8]) -> Option { + pub fn write(&mut self, index: u16, name: &[u8], key: &[u8]) -> Result { 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) + + self.mem.write(index_name, name).map_err(|_| ())?; + self.mem.write(index_key, key).map_err(|_| ())?; + + Ok(index) } pub fn delete(&mut self, index: u16) -> Option { @@ -116,7 +108,7 @@ impl<'a> Tokens<'a> { Some(index) } - pub fn wipe_all_tokens(&mut self) -> u8 { + pub fn soft_wipe_all_tokens(&mut self) -> u8 { let mut inc = 0; for n in 0..self.capacity { @@ -128,4 +120,8 @@ impl<'a> Tokens<'a> { } inc } + + pub fn hard_wipe_all_tokens(&mut self) { + self.mem.erase(0, self.mem.capacity()).unwrap(); + } } diff --git a/tools/menu.py b/tools/menu.py index 71510e4..0cea0fb 100644 --- a/tools/menu.py +++ b/tools/menu.py @@ -37,18 +37,6 @@ def process_secret(secret: str) -> bytes: 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() @@ -98,12 +86,12 @@ def main(argv: list[str]): 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) + conn.write(date_items["year"]) + conn.write(date_items["month"]) + conn.write(date_items["day"]) + conn.write(date_items["hours"]) + conn.write(date_items["minutes"]) + conn.write(date_items["seconds"]) resp = conn.read() if resp != OK: print(f"Error in the communication: Error {resp}")