otp_offline 0.2.0

//! A simple key-value store for OTP (One-Time Password) data to prevent replay attacks.
//!
//! This module provides a `Store` that persists OTP data to a file.
//! It supports loading, updating, and storing OTP entries, with a lookup table
//! for access by public ID and preventing replay attacks.
//!
//! The file format is a text-based key-value format, where each line corresponds
//! to an OTP entry. Each entry includes:
//! - A public ID (encoded in modhex),
//! - Usage counter, session counter, private ID and the secret key.
//!
//! The store is thread-safe for updates via internal locking.
//!
//! The rows are fixed-length to allow efficient in-place updates without rewriting
//! the entire file. Only modified entries are written back to disk based on the cached
//! data in memory.
//!
//! <div class="warning">
//!
//! **Therefore, it is important to not modify the file while the program
//! is running, as this may corrupt the data.**
//!
//! </div>
//!
//! ```text
//! ######################### Do not modify this file or this line while the program is running this will result in a broken file #########################
//! cbcdcecfcgch = usage_counter:00000 session_counter:000 private_id:0708090a0b0c key:00000000000000000000000000000000
//! ```
//!
//! # Example
//!
//! ```
//! use otp_offline::simple_store::{Store, Data};
//! use otp_offline::otp::{DecryptedOtp, PublicId, PrivateId, DecryptedPrivateData};
//!
//! let mut store = Store::create("otp_data").unwrap();
//! let otp = DecryptedOtp {
//!     id: PublicId{raw_bytes: [1, 2, 3, 4, 5, 6]},
//!     private: DecryptedPrivateData {
//!         id: PrivateId{raw_bytes: [7, 8, 9, 10, 11, 12]},
//!         usage_counter: 0,
//!         session_counter: 0,
//!         timestamp: 0,
//!         random: [0; 2],
//!     },
//! };
//! store.provision_new_otp(&otp, &[0; 16]);
//! ```

use std::{
    collections::HashMap,
    default::Default,
    fs::File,
    io::{Read, Seek, SeekFrom, Write},
    sync::Mutex,
    time::SystemTime,
};

use crate::{
    modhex::ModHex,
    otp::{self, DecryptedOtp, DecryptedPrivateData, Otp, PrivateId, PublicId},
    store::{OtpStore, StoreError},
};

/// A store for OTP data persisted in a file.
///
/// The store maintains an in-memory representation of the data for fast access,
/// and writes changes back to the file when updated.
pub struct Store {
    /// Path to the file
    file_path: String,

    /// In-memory entries of OTP data.
    entries: Vec<Data>,

    /// Lookup table mapping public IDs to entry indices.
    lookup: HashMap<PublicId, usize>,

    /// Mutex to ensure thread-safe writes to the file.
    write_lock: Mutex<()>,
}

/// Data representing an OTP entry in the store.
pub struct Data {
    /// The encryption key used for this OTP entry.
    pub key: [u8; 16],

    /// Previous private data associated with the OTP.
    pub previous: DecryptedPrivateData,
}

#[derive(Debug)]
pub enum Error {
    /// An I/O error occurred while reading or writing the file.
    IoError(std::io::Error),

    /// An error occurred while parsing the file.
    Parsing,
}

// Calculate the line length (all lines have the same fixed length)
// Each line is LINE_LENGTH bytes + 1 for newline
const LINE_LENGTH_WITH_NEWLINE: usize = 90 + 1;

impl std::error::Error for Error {}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Error::IoError(error) => write!(f, "{error}"),
            Error::Parsing => write!(f, "Error parsing file"),
        }
    }
}

impl Default for Data {
    fn default() -> Self {
        Self {
            key: [0; 16],
            previous: DecryptedPrivateData {
                id: PrivateId { raw_bytes: [0; 6] },
                usage_counter: 0,
                session_counter: 0,
                timestamp: 0,
                random: [0; 2],
            },
        }
    }
}

impl Store {
    /// Loads the store from a file.
    ///
    /// If the file does not exist, an empty store is returned.
    ///
    /// # Errors
    /// Returns an error if the file cannot be opened or parsed.
    pub fn create(file_path: &str) -> Result<Self, Error> {
        let mut f = match File::open(file_path) {
            Ok(file) => file,
            Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
                // File does not exist, create a new empty store and write header
                let store = Store {
                    file_path: String::from(file_path),
                    lookup: HashMap::new(),
                    entries: Vec::new(),
                    write_lock: Mutex::new(()),
                };
                store.store_all().map_err(Error::IoError)?;
                return Ok(store);
            }
            Err(e) => return Err(Error::IoError(e)),
        };

        let mut buffer = String::new();
        f.read_to_string(&mut buffer).map_err(Error::IoError)?;

        let mut lookup = HashMap::new();
        let mut entries = Vec::new();

        // In case someone messed up the line endings
        let rewrite_whole_file = buffer.contains("\r\n");

        for (line_index, line) in buffer.trim().lines().skip(1).enumerate() {
            let line = line.trim();
            if line.is_empty() {
                continue;
            }

            let Some((public_id_str, data)) = line.split_once('=') else {
                return Err(Error::Parsing);
            };

            let mut public_id = PublicId { raw_bytes: [0; 6] };
            public_id.raw_bytes.copy_from_slice(
                ModHex::try_from(public_id_str.trim())
                    .map_err(|_| Error::Parsing)?
                    .raw_bytes(),
            );

            // Parse key-value pairs flexibly
            let mut usage_counter: Option<u16> = None;
            let mut session_counter: Option<u8> = None;
            let mut private_id: Option<[u8; 6]> = None;
            let mut key: Option<[u8; 16]> = None;

            for part in data.split_whitespace() {
                if let Some(value) = part.strip_prefix("u_cnt:") {
                    usage_counter = Some(value.parse().map_err(|_| Error::Parsing)?);
                } else if let Some(value) = part.strip_prefix("s_cnt:") {
                    session_counter = Some(value.parse().map_err(|_| Error::Parsing)?);
                } else if let Some(value) = part.strip_prefix("pid:") {
                    private_id = Some(
                        decode_hex(value)
                            .map_err(|_| Error::Parsing)?
                            .try_into()
                            .map_err(|_| Error::Parsing)?,
                    );
                } else if let Some(value) = part.strip_prefix("key:") {
                    key = Some(
                        decode_hex(value)
                            .map_err(|_| Error::Parsing)?
                            .try_into()
                            .map_err(|_| Error::Parsing)?,
                    );
                }
            }

            // Ensure all required fields are present
            let usage_counter = usage_counter.ok_or(Error::Parsing)?;
            let session_counter = session_counter.ok_or(Error::Parsing)?;
            let private_id = private_id.ok_or(Error::Parsing)?;
            let key = key.ok_or(Error::Parsing)?;

            let previous = DecryptedPrivateData {
                id: crate::otp::PrivateId {
                    raw_bytes: private_id,
                },
                usage_counter,
                session_counter,
                timestamp: 0,
                random: [0; 2],
            };

            entries.push(Data { key, previous });
            lookup.insert(public_id, line_index);
        }

        let store = Store {
            file_path: String::from(file_path),
            lookup,
            entries,
            write_lock: Mutex::new(()),
        };

        // Rewrite the whole file to fix line endings
        if rewrite_whole_file {
            store.store_all().map_err(Error::IoError)?;
        }

        Ok(store)
    }

    /// Retrieves an entry by its public ID.
    ///
    /// Returns `None` if the ID is not found.
    pub fn get(&self, id: PublicId) -> Option<&Data> {
        self.lookup.get(&id).map(|&idx| &self.entries[idx])
    }

    /// Validates an OTP string against the store for a known OTP.
    ///
    /// If valid, updates the store with the new OTP data.
    ///
    /// # Errors
    /// If the OTP is invalid, unknown, or cannot be decrypted.
    pub fn validate(&mut self, otp_str: &str) -> Result<(), StoreError> {
        let otp = match Otp::from_modhex(otp_str) {
            Ok(otp) => otp,
            Err(err) => {
                return Err(StoreError::Otp(err));
            }
        };

        let Some(previous) = self.get(otp.id) else {
            return Err(StoreError::UnknownPublicId);
        };

        // Try decryption
        let decrypted_otp = match otp.decrypt(&previous.key) {
            Ok(otp) => otp,
            Err(err) => {
                return Err(StoreError::Otp(err));
            }
        };

        match decrypted_otp.validate(&otp::DecryptedOtp {
            id: decrypted_otp.id,
            private: previous.previous.clone(),
        }) {
            Ok(()) => self
                .update(&decrypted_otp, &previous.key.clone())
                .map_err(|e| StoreError::Other(Box::new(e))),
            Err(err) => Err(StoreError::Validation(err)),
        }
    }

    /// Adds a new trusted OTP entry to the store.
    ///
    /// This should only be used for adding new OTPs in the store.
    /// In all other cases, use the [`validate`](Self::validate) method to ensure proper validation and updating.
    ///
    /// # Errors
    /// If an I/O error occurs while writing to the file.
    pub fn provision_new_otp(&mut self, otp: &DecryptedOtp, key: &[u8]) -> Result<(), Error> {
        self.update(otp, key)
    }

    /// Updates the store with new OTP data.
    ///
    /// If the OTP already exists, it is updated. Otherwise, it is added and the store is persisted to disk.
    ///
    /// # Errors
    /// If an I/O error occurs while writing to the file.
    fn update(&mut self, otp: &DecryptedOtp, key: &[u8]) -> Result<(), Error> {
        let mut is_new = false;

        if let Some(existing_index) = self.lookup.get(&otp.id) {
            self.entries[*existing_index].previous = otp.private.clone();
        } else {
            let mut new_data = Data {
                previous: otp.private.clone(),
                ..Default::default()
            };
            new_data.key.copy_from_slice(key);

            self.entries.push(new_data);
            self.lookup.insert(otp.id, self.entries.len() - 1);
            is_new = true;
        }

        self.store(otp.id, is_new).map_err(Error::IoError)
    }

    /// Formats a single line for storage in the file.
    fn write_otp<W: std::io::Write>(
        writer: &mut W,
        public_id: PublicId,
        entry: &Data,
        write_provision_data: bool,
    ) -> Result<(), std::io::Error> {
        write!(
            writer,
            "{} = u_cnt:{:05}",
            ModHex::from(&public_id.raw_bytes[..]),
            entry.previous.usage_counter
        )?;
        write!(writer, " s_cnt:{:03}", entry.previous.session_counter)?;

        if write_provision_data {
            write!(writer, " pid:")?;
            for b in &entry.previous.id.raw_bytes {
                write!(writer, "{b:02x}")?;
            }
            write!(writer, " key:")?;
            for b in entry.key {
                write!(writer, "{b:02x}")?;
            }
        }

        Ok(())
    }

    /// Stores the entry identified by `public_id` to disk.
    ///
    /// If the file doesn't exist, all entries are written.
    /// If the entry is new, it is appended.
    /// Otherwise, the existing line is replaced.
    fn store(&mut self, public_id: PublicId, is_new: bool) -> Result<(), std::io::Error> {
        let _lock = self.write_lock.lock().unwrap();

        if !std::path::Path::new(&self.file_path).exists() {
            // File doesn't exist, write everything
            return self.store_all();
        }

        // Open file for writing
        let mut file = std::fs::OpenOptions::new()
            .write(true)
            .open(&self.file_path)?;

        let line_index = *self.lookup.get(&public_id).expect("Id must be present");
        let entry = &self.entries[line_index];

        if is_new {
            // Entry doesn't exist yet, append to end
            file.seek(SeekFrom::End(0))?;
            writeln!(file)?;
        } else {
            // Seek to the correct position to replace the line
            let position = (line_index + 1) * LINE_LENGTH_WITH_NEWLINE;
            file.seek(SeekFrom::Start(position as u64))?;
        }

        Self::write_otp(&mut file, public_id, entry, is_new)
    }

    /// Writes all entries to the file.
    ///
    /// Used when creating a new file or replacing an existing one.
    fn store_all(&self) -> Result<(), std::io::Error> {
        let _lock = self.write_lock.lock().unwrap();

        let mut file = File::create(&self.file_path)?;
        let mut first = true;

        write!(
            file,
            "# Do not modify this file while the program is running this will result in a broken file #"
        )?;

        for (&public_id, &index) in &self.lookup {
            if !first {
                writeln!(file)?;
            }
            first = false;

            Self::write_otp(&mut file, public_id, &self.entries[index], true)
                .map_err(|_| std::io::Error::other("Format error"))?;
        }
        Ok(())
    }
}

/// Helper function to parse hex string to byte array
pub(crate) fn decode_hex(s: &str) -> Result<Vec<u8>, std::num::ParseIntError> {
    (0..s.len())
        .step_by(2)
        .map(|i| u8::from_str_radix(&s[i..i + 2], 16))
        .collect()
}

impl OtpStore for Store {
    fn validate(&mut self, otp_str: &str, _now: SystemTime) -> Result<(), StoreError> {
        self.validate(otp_str)
    }
}