wallet-dat 0.1.0

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Copyright (c) DUSK NETWORK. All rights reserved.

//!
//! The wallet.dat structure is a binary file containing the seed we get from a bip39 mnemonic,
//! in encrypted form via a password
#![no_std]
extern crate alloc;

use constants::*;
use crypto::{decrypt, encrypt};
use error::{Result, WalletDataFileError as Error};

use seed::Seed;

use alloc::vec::Vec;

mod crypto;
mod error;
mod seed;
mod constants {
    /// Binary prefix for new binary file format
    pub const MAGIC: u32 = 0x72736b;
    /// Binary prefix for old Dusk wallet files
    pub const OLD_MAGIC: u32 = 0x1d0c15;
    pub const FILE_TYPE: u16 = 0x0200;
    pub const RESERVED: u16 = 0x0000;
    pub const LATEST_VERSION: super::Version = (0, 0, 2, 0, false);
}

type Version = (u8, u8, u8, u8, bool);

/// Versions of the potential wallet DAT files we read
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum DatFileVersion {
    /// Legacy the oldest format
    Legacy,
    /// Preciding legacy, we have the old one
    OldWalletCli(Version),
    /// The newest one. All new saves are saved in this file format
    RuskBinaryFileFormat(Version),
}

impl DatFileVersion {
    /// Checks if the file version is older than the latest Rust Binary file
    /// format
    pub fn is_old(&self) -> bool {
        matches!(self, Self::Legacy | Self::OldWalletCli(_))
    }
}

/// Make sense of the payload and return it
pub fn get_seed_and_address(
    file: DatFileVersion,
    mut bytes: Vec<u8>,
    pwd: &str,
) -> Result<(Seed, u64)> {
    let pwd = blake3::hash(pwd.as_bytes());
    let pwd = pwd.as_bytes();

    match file {
        DatFileVersion::Legacy => {
            if bytes[1] == 0 && bytes[2] == 0 {
                bytes.drain(..3);
            }

            bytes = decrypt(&bytes, pwd)?;

            // get our seed
            let seed = Seed::from_reader(&bytes[..]).map_err(|_| Error::WalletFileCorrupted)?;

            Ok((seed, 1))
        }
        DatFileVersion::OldWalletCli((major, minor, _, _, _)) => {
            bytes.drain(..5);

            let result: Result<(Seed, _)> = match (major, minor) {
                (1, 0) => {
                    let content = decrypt(&bytes, pwd)?;
                    let buff = &content[..];

                    let seed = Seed::from_reader(buff).map_err(|_| Error::WalletFileCorrupted)?;

                    Ok((seed, 0))
                }
                (2, 0) => {
                    let content = decrypt(&bytes, pwd)?;
                    let buff = &content[..];

                    // extract seed
                    let seed = Seed::from_reader(buff).map_err(|_| Error::WalletFileCorrupted)?;

                    // extract addresses count
                    Ok((seed, 0))
                }
                _ => Err(Error::UnknownFileVersion(major, minor)),
            };

            result
        }
        DatFileVersion::RuskBinaryFileFormat(version) => {
            let rest = bytes.get(12..(12 + 96));

            if let Some(rest) = rest {
                let content = decrypt(rest, pwd)?;

                if let Some(seed_buff) = content.get(0..65) {
                    // first 64 bytes are the seed
                    let seed = Seed::from_reader(&seed_buff[0..64])
                        .map_err(|_| Error::WalletFileCorrupted)?;

                    match version {
                        (0, 0, 2, 0, false) => {
                            if let Some(last_pos_bytes) = content.get(64..72) {
                                let last_pos = match last_pos_bytes.try_into() {
                                    Ok(last_pos_bytes) => u64::from_le_bytes(last_pos_bytes),
                                    Err(_) => return Err(Error::NoLastPosFound),
                                };

                                Ok((seed, last_pos))
                            } else {
                                Err(Error::WalletFileCorrupted)
                            }
                        }
                        _ => Ok((seed, 0)),
                    }
                } else {
                    Err(Error::WalletFileCorrupted)
                }
            } else {
                Err(Error::WalletFileCorrupted)
            }
        }
    }
}

/// From the first 12 bytes of the file (header), we check version
///
/// https://github.com/dusk-network/rusk/wiki/Binary-File-Format/#header
pub fn check_version(bytes: Option<&[u8]>) -> Result<DatFileVersion> {
    match bytes {
        Some(bytes) => {
            let header_bytes: [u8; 4] = bytes[0..4]
                .try_into()
                .map_err(|_| Error::WalletFileCorrupted)?;

            let magic = u32::from_le_bytes(header_bytes) & 0x00ffffff;

            if magic == OLD_MAGIC {
                // check for version information
                let (major, minor) = (bytes[3], bytes[4]);

                Ok(DatFileVersion::OldWalletCli((major, minor, 0, 0, false)))
            } else {
                let header_bytes = bytes[0..8]
                    .try_into()
                    .map_err(|_| Error::WalletFileCorrupted)?;

                let number = u64::from_be_bytes(header_bytes);

                let magic_num = (number & 0xFFFFFF00000000) >> 32;

                if (magic_num as u32) != MAGIC {
                    return Ok(DatFileVersion::Legacy);
                }

                let file_type = (number & 0x000000FFFF0000) >> 16;
                let reserved = number & 0x0000000000FFFF;

                if file_type != FILE_TYPE as u64 {
                    return Err(Error::WalletFileCorrupted);
                };

                if reserved != RESERVED as u64 {
                    return Err(Error::WalletFileCorrupted);
                };

                let version_bytes = bytes[8..12]
                    .try_into()
                    .map_err(|_| Error::WalletFileCorrupted)?;

                let version = u32::from_be_bytes(version_bytes);

                let major = (version & 0xff000000) >> 24;
                let minor = (version & 0x00ff0000) >> 16;
                let patch = (version & 0x0000ff00) >> 8;
                let pre = (version & 0x000000f0) >> 4;
                let higher = version & 0x0000000f;

                let pre_higher = matches!(higher, 1);

                Ok(DatFileVersion::RuskBinaryFileFormat((
                    major as u8,
                    minor as u8,
                    patch as u8,
                    pre as u8,
                    pre_higher,
                )))
            }
        }
        None => Err(Error::WalletFileCorrupted),
    }
}

pub fn encrypt_seed(seed: &[u8; 64], pwd: &str, last_pos: u64) -> Result<Vec<u8>> {
    let mut header = Vec::with_capacity(12);
    header.extend_from_slice(&MAGIC.to_be_bytes());
    // File type = Rusk Wallet (0x02)
    header.extend_from_slice(&FILE_TYPE.to_be_bytes());
    // Reserved (0x0)
    header.extend_from_slice(&RESERVED.to_be_bytes());
    // Version
    header.extend_from_slice(&version_bytes(LATEST_VERSION));

    let mut payload = Vec::from(seed);

    payload.extend(last_pos.to_le_bytes());

    let pwd = blake3::hash(pwd.as_bytes());
    let pwd = pwd.as_bytes();

    // encrypt the payload
    payload = encrypt(&payload, pwd)?;

    let mut content = Vec::with_capacity(header.len() + payload.len());

    content.extend_from_slice(&header);
    content.extend_from_slice(&payload);

    Ok(content)
}

pub(crate) fn version_bytes(version: Version) -> [u8; 4] {
    u32::from_be_bytes([version.0, version.1, version.2, version.3]).to_be_bytes()
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloc::vec::Vec;

    #[test]
    fn distiction_between_versions() {
        // with magic number
        let old_wallet_file = Vec::from([0x15, 0x0c, 0x1d, 0x02, 0x00]);
        // no magic number just nonsense bytes
        let legacy_file = Vec::from([
            0xab, 0x38, 0x81, 0x3b, 0xfc, 0x79, 0x11, 0xf9, 0x86, 0xd6, 0xd0,
        ]);
        // new header
        let new_file = Vec::from([
            0x00, 0x72, 0x73, 0x6b, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00,
        ]);

        assert_eq!(
            check_version(Some(&old_wallet_file)).unwrap(),
            DatFileVersion::OldWalletCli((2, 0, 0, 0, false))
        );

        assert_eq!(
            check_version(Some(&legacy_file)).unwrap(),
            DatFileVersion::Legacy
        );

        assert_eq!(
            check_version(Some(&new_file)).unwrap(),
            DatFileVersion::RuskBinaryFileFormat((0, 0, 1, 0, false))
        );
    }

    #[test]
    fn generate_latest_version() {
        let seed: [u8; 64] = [0; 64];
        let encryped = encrypt_seed(&seed, "password", 304).unwrap();

        let version = check_version(Some(&encryped)).unwrap();

        assert_eq!(
            version,
            DatFileVersion::RuskBinaryFileFormat(LATEST_VERSION)
        );

        let (returned_seed, last_pos) =
            get_seed_and_address(version, encryped, "password").unwrap();

        assert_eq!(seed, returned_seed.as_bytes());

        assert_eq!(last_pos, 304);
    }
}