use crate::{
err::ForensicError,
traits::registry::{RegHiveKey, RegValue, RegistryKeyInfo, RegistryReader},
};
use std::{cell::RefCell, collections::BTreeMap};
use super::time::Filetime;
#[derive(Clone, Debug)]
pub struct TestingRegistry {
pub cell: BTreeMap<String, MountedCell>,
pub cached: RefCell<BTreeMap<RegHiveKey, String>>,
pub counter: RefCell<isize>,
}
impl Default for TestingRegistry {
fn default() -> Self {
Self::new()
}
}
impl TestingRegistry {
pub fn empty() -> Self {
Self {
cell: BTreeMap::new(),
cached: RefCell::new(basic_cache()),
counter: RefCell::default(),
}
}
pub fn new() -> Self {
Self {
cell: basic_registry(),
cached: RefCell::new(basic_cache()),
counter: RefCell::new(0),
}
}
pub fn increase_counter(&self) -> isize {
let mut borrowed = self.counter.borrow_mut();
let ret = *borrowed;
*borrowed += 1;
ret
}
pub fn add_value(&mut self, path: &str, value: &str, data: RegValue) {
let (hkey, rest) = match path.split_once(|v| v == '/' || v == '\\') {
Some(v) => v,
None => {
return self
.cell
.entry(path.to_string())
.or_insert(MountedCell::new(path))
.add_value("", value, data)
}
};
self.cell
.entry(hkey.to_string())
.or_insert(MountedCell::new(hkey))
.add_value(rest, value, data);
}
pub fn contains(&self, path: &str) -> bool {
let (hkey, rest) = match path.split_once(|v| v == '/' || v == '\\') {
Some(v) => v,
None => return self.cell.contains_key(path),
};
let hive = match self.cell.get(hkey) {
Some(v) => v,
None => return false,
};
hive.contains_key(rest)
}
pub fn get_value(&self, path: &str, value: &str) -> Option<RegValue> {
let (hkey, rest) = match path.split_once(|v| v == '/' || v == '\\') {
Some(v) => v,
None => (path, ""),
};
let hive = match self.cell.get(hkey) {
Some(v) => v,
None => return None,
};
hive.get_value(rest, value)
}
pub fn get_values(&self, path: &str) -> Option<Vec<String>> {
let (hkey, rest) = match path.split_once(|v| v == '/' || v == '\\') {
Some(v) => v,
None => (path, ""),
};
let hive = match self.cell.get(hkey) {
Some(v) => v,
None => return None,
};
Some(hive.get_values(rest))
}
pub fn get_keys(&self, path: &str) -> Option<Vec<String>> {
let (hkey, rest) = match path.split_once(|v| v == '/' || v == '\\') {
Some(v) => v,
None => (path, ""),
};
let hive = match self.cell.get(hkey) {
Some(v) => v,
None => return None,
};
Some(hive.get_keys(rest))
}
}
#[derive(Clone, Debug, Default)]
pub struct MountedCell {
pub name: String,
pub keys: BTreeMap<String, MountedCell>,
pub values: BTreeMap<String, RegValue>,
}
impl MountedCell {
pub fn new(name: &str) -> Self {
Self {
name: name.into(),
keys: BTreeMap::new(),
values: BTreeMap::new(),
}
}
pub fn add_key(&mut self, path: &str) {
if path.is_empty() {
return;
}
let (first, rest) = match path.split_once(|v| v == '/' || v == '\\') {
Some(v) => v,
None => {
self.keys
.entry(path.to_string())
.or_insert(MountedCell::new(path))
.add_key(path);
return;
}
};
self.keys
.entry(first.to_string())
.or_insert(MountedCell::new(first))
.add_key(rest);
}
pub fn contains_key(&self, path: &str) -> bool {
let (first, rest) = match path.split_once(|v| v == '/' || v == '\\') {
Some(v) => v,
None => return self.keys.contains_key(path),
};
let hive = match self.keys.get(first) {
Some(v) => v,
None => return false,
};
hive.contains_key(rest)
}
pub fn add_value(&mut self, path: &str, value: &str, data: RegValue) {
if path.is_empty() {
self.values.insert(value.into(), data);
return;
}
let (first, rest) = match path.split_once(|v| v == '/' || v == '\\') {
Some(v) => v,
None => {
self.keys
.entry(path.to_string())
.or_insert(MountedCell::new(path))
.add_value("", value, data);
return;
}
};
self.keys
.entry(first.to_string())
.or_insert(MountedCell::new(first))
.add_value(rest, value, data);
}
pub fn get_value(&self, path: &str, value: &str) -> Option<RegValue> {
if path.is_empty() {
return self.values.get(value).cloned();
}
let (first, rest) = match path.split_once(|v| v == '/' || v == '\\') {
Some(v) => v,
None => return self.keys.get(path)?.get_value("", value),
};
self.keys.get(first)?.get_value(rest, value)
}
pub fn get_values(&self, path: &str) -> Vec<String> {
if path.is_empty() {
return self
.values
.keys()
.map(|v| v.to_string())
.collect();
}
let (first, rest) = match path.split_once(|v| v == '/' || v == '\\') {
Some(v) => v,
None => {
return match self.keys.get(path) {
Some(v) => v.get_values(""),
None => Vec::new(),
}
}
};
match self.keys.get(first) {
Some(v) => v.get_values(rest),
None => Vec::new(),
}
}
pub fn get_keys(&self, path: &str) -> Vec<String> {
if path.is_empty() {
return self
.keys
.keys()
.map(|v| v.to_string())
.collect();
}
let (first, rest) = match path.split_once(|v| v == '/' || v == '\\') {
Some(v) => v,
None => {
return match self.keys.get(path) {
Some(v) => v.get_keys(""),
None => Vec::new(),
}
}
};
match self.keys.get(first) {
Some(v) => v.get_keys(rest),
None => Vec::new(),
}
}
}
impl RegistryReader for TestingRegistry {
fn from_file(
&self,
_file: Box<dyn crate::traits::vfs::VirtualFile>,
) -> crate::err::ForensicResult<Box<dyn RegistryReader>> {
Ok(Box::new(TestingRegistry::new()))
}
fn from_fs(
&self,
_fs: Box<dyn crate::traits::vfs::VirtualFileSystem>,
) -> crate::err::ForensicResult<Box<dyn RegistryReader>> {
Ok(Box::new(TestingRegistry::new()))
}
fn open_key(
&self,
hkey: crate::traits::registry::RegHiveKey,
key_name: &str,
) -> crate::err::ForensicResult<crate::traits::registry::RegHiveKey> {
let mut borrowed = self.cached.borrow_mut();
let (hkey, path) = match borrowed.get(&hkey) {
Some(v) => {
let full_path = format!("{}\\{}", v, key_name);
if !self.contains(&full_path) {
return Err(ForensicError::missing_string(format!(
"Key path {} not found",
full_path
)));
}
let handle = self.increase_counter();
(handle, full_path)
}
None => return Err(ForensicError::missing_str("Hkey not found")),
};
borrowed.insert(RegHiveKey::Hkey(hkey), path);
Ok(RegHiveKey::Hkey(hkey))
}
fn read_value(
&self,
hkey: crate::traits::registry::RegHiveKey,
value_name: &str,
) -> crate::err::ForensicResult<RegValue> {
let borrowed = self.cached.borrow();
let key_path = borrowed
.get(&hkey)
.ok_or_else(|| ForensicError::missing_str("HKey not found"))?;
let value = self.get_value(key_path, value_name).ok_or_else(|| {
ForensicError::missing_string(format!("Value {}\\{} not found", key_path, value_name))
})?;
Ok(value)
}
fn enumerate_values(
&self,
hkey: crate::traits::registry::RegHiveKey,
) -> crate::err::ForensicResult<Vec<String>> {
let borrowed = self.cached.borrow();
let key_path = borrowed
.get(&hkey)
.ok_or_else(|| ForensicError::missing_str("HKey not found"))?;
let value = self.get_values(key_path).ok_or_else(|| {
ForensicError::missing_string(format!("Values for {} not found", key_path))
})?;
Ok(value)
}
fn enumerate_keys(
&self,
hkey: crate::traits::registry::RegHiveKey,
) -> crate::err::ForensicResult<Vec<String>> {
let borrowed = self.cached.borrow();
let key_path = borrowed
.get(&hkey)
.ok_or_else(|| ForensicError::missing_str("HKey not found"))?;
let value = self.get_keys(key_path).ok_or_else(|| {
ForensicError::missing_string(format!("Keys for {} not found", key_path))
})?;
Ok(value)
}
fn key_at(
&self,
hkey: crate::traits::registry::RegHiveKey,
pos: u32,
) -> crate::err::ForensicResult<String> {
let borrowed = self.cached.borrow();
let key_path = borrowed
.get(&hkey)
.ok_or_else(|| ForensicError::missing_str("HKey not found"))?;
let mut value = self.get_keys(key_path).ok_or_else(|| {
ForensicError::missing_string(format!("Keys for {} not found", key_path))
})?;
let pos = pos as usize;
if pos > value.len() {
return Err(ForensicError::NoMoreData);
}
Ok(value.remove(pos))
}
fn value_at(
&self,
hkey: crate::traits::registry::RegHiveKey,
pos: u32,
) -> crate::err::ForensicResult<String> {
let borrowed = self.cached.borrow();
let key_path = borrowed
.get(&hkey)
.ok_or_else(|| ForensicError::missing_str("HKey not found"))?;
let mut value = self.get_values(key_path).ok_or_else(|| {
ForensicError::missing_string(format!("Values for {} not found", key_path))
})?;
let pos = pos as usize;
if pos > value.len() {
return Err(ForensicError::NoMoreData);
}
Ok(value.remove(pos))
}
fn key_info(&self, hkey: RegHiveKey) -> crate::err::ForensicResult<crate::traits::registry::RegistryKeyInfo> {
let borrowed = self.cached.borrow();
let key_path = borrowed
.get(&hkey)
.ok_or_else(|| ForensicError::missing_str("HKey not found"))?;
let value = self.get_values(key_path).ok_or_else(|| {
ForensicError::missing_string(format!("Values for {} not found", key_path))
})?;
let keys = self.get_keys(key_path).ok_or_else(|| {
ForensicError::missing_string(format!("Values for {} not found", key_path))
})?;
Ok(RegistryKeyInfo {
last_write_time : Filetime::new(0),
subkeys : keys.len() as u32,
values : value.len() as u32,
max_subkey_name_length : keys.iter().map(|v| v.len()).fold(0, |acc, e| e.max(acc)) as u32,
max_value_name_length: value.iter().map(|v| v.len()).fold(0, |acc, e| e.max(acc)) as u32,
max_value_length: 0,
})
}
}
fn basic_cache() -> BTreeMap<RegHiveKey, String> {
{
let mut map = BTreeMap::new();
for (k, p) in [
(RegHiveKey::HkeyLocalMachine, "HKLM"),
(RegHiveKey::HkeyCurrentUser, "HKCU"),
(RegHiveKey::HkeyUsers, "HKU"),
(RegHiveKey::HkeyClassesRoot, "HKCR"),
] {
map.insert(k, p.to_string());
}
map
}
}
fn basic_registry() -> BTreeMap<String, MountedCell> {
let mut map = BTreeMap::new();
for k in ["HKLM", "HKCU", "HKCR"] {
map.insert(k.to_string(), MountedCell::new(k));
}
let mut hkcu_cell = MountedCell::new("HKU");
hkcu_cell.add_value(
"S-1-5-21-1366093794-4292800403-1155380978-513\\Volatile Environment",
"USERPROFILE",
RegValue::from_str(r"C:\Users\Tester"),
);
hkcu_cell.add_value(
"S-1-5-21-1366093794-4292800403-1155380978-513\\Volatile Environment",
"APPDATA",
RegValue::from_str(r"C:\Users\Tester\AppData\Roaming"),
);
hkcu_cell.add_value(
"S-1-5-21-1366093794-4292800403-1155380978-513\\Volatile Environment",
"LOCALAPPDATA",
RegValue::from_str(r"C:\Users\Tester\AppData\Local"),
);
hkcu_cell.add_value(
"S-1-5-21-1366093794-4292800403-1155380978-513\\Volatile Environment",
"USERDOMAIN",
RegValue::from_str(r"TestMachine"),
);
hkcu_cell.add_value(
"S-1-5-21-1366093794-4292800403-1155380978-513\\Volatile Environment",
"USERNAME",
RegValue::from_str(r"Tester"),
);
map.insert("HKU".into(), hkcu_cell);
map
}
pub fn init_testing_logger() {
let rcv = crate::notifications::testing_notifier_dummy();
std::thread::spawn(move || loop {
let msg = match rcv.recv() {
Ok(v) => v,
Err(_) => return,
};
println!(
"{:?} - {} - {}:{} - {}",
msg.r#type, msg.module, msg.file, msg.line, msg.data
);
});
let rcv = crate::logging::testing_logger_dummy();
std::thread::spawn(move || loop {
let msg = match rcv.recv() {
Ok(v) => v,
Err(_) => return,
};
println!(
"{:?} - {} - {}:{} - {}",
msg.level, msg.module, msg.file, msg.line, msg.data
);
});
}