Mirror/rayhunter
1
0
Fork 0
mirror of https://github.com/EFForg/rayhunter.git synced 2026-04-30 01:19:58 -07:00
Code Issues Packages Projects Releases Wiki Activity

daemon: switch to writing heuristics output to ND-JSON

Browse Source 
ND-JSON (newline-delimited JSON) is just a file with a list of JSON
objects separated by newlines. This way, as the analyzer harness
processes new packets, it can simply append JSON-serialized results
to a file without parsing the entire thing first.

Also simplifies the analysis stuff to all operate in the diag thread.
This commit is contained in:
Will Greenberg
2024-05-09 14:46:41 -07:00
parent 4a5bede4ee
commit bfc688ad21
10 changed files with 256 additions and 224 deletions
Download Patch File Download Diff File Expand all files Collapse all files

139
bin/src/qmdl_store.rs
View File

@@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize};
use chrono::{DateTime, Local};
#[derive(Debug, Error)]
pub enum QmdlStoreError {
pub enum RecordingStoreError {
#[error("Can't close an entry when there's no current entry")]
NoCurrentEntry,
#[error("Couldn't create file: {0}")]
@@ -22,7 +22,7 @@ pub enum QmdlStoreError {
ParseManifestError(toml::de::Error)
}
pub struct QmdlStore {
pub struct RecordingStore {
pub path: PathBuf,
pub manifest: Manifest,
pub current_entry: Option<usize>, // index into manifest
@@ -38,7 +38,8 @@ pub struct ManifestEntry {
pub name: String,
pub start_time: DateTime<Local>,
pub last_message_time: Option<DateTime<Local>>,
pub size_bytes: usize,
pub qmdl_size_bytes: usize,
pub analysis_size_bytes: usize,
}
impl ManifestEntry {
@@ -48,113 +49,142 @@ impl ManifestEntry {
name: format!("{}", now.timestamp()),
start_time: now,
last_message_time: None,
size_bytes: 0,
qmdl_size_bytes: 0,
analysis_size_bytes: 0,
}
}
pub fn get_qmdl_filepath<P: AsRef<Path>>(&self, path: P) -> PathBuf {
let mut filepath = path.as_ref().join(&self.name);
filepath.set_extension("qmdl");
filepath
}
pub fn get_analysis_filepath<P: AsRef<Path>>(&self, path: P) -> PathBuf {
let mut filepath = path.as_ref().join(&self.name);
filepath.set_extension("ndjson");
filepath
}
}
impl QmdlStore {
impl RecordingStore {
// Returns whether a directory with a "manifest.toml" exists at the given
// path (though doesn't check if that manifest is valid)
pub async fn exists<P>(path: P) -> Result<bool, QmdlStoreError> where P: AsRef<Path> {
pub async fn exists<P>(path: P) -> Result<bool, RecordingStoreError> where P: AsRef<Path> {
let manifest_path = path.as_ref().join("manifest.toml");
let dir_exists = try_exists(path).await.map_err(QmdlStoreError::OpenDirError)?;
let manifest_exists = try_exists(manifest_path).await.map_err(QmdlStoreError::ReadManifestError)?;
let dir_exists = try_exists(path).await.map_err(RecordingStoreError::OpenDirError)?;
let manifest_exists = try_exists(manifest_path).await.map_err(RecordingStoreError::ReadManifestError)?;
Ok(dir_exists && manifest_exists)
}
// Loads an existing QmdlStore at the given path. Errors if no store exists,
// Loads an existing RecordingStore at the given path. Errors if no store exists,
// or if it's malformed.
pub async fn load<P>(path: P) -> Result<Self, QmdlStoreError> where P: AsRef<Path> {
pub async fn load<P>(path: P) -> Result<Self, RecordingStoreError> where P: AsRef<Path> {
let path: PathBuf = path.as_ref().to_path_buf();
let manifest = QmdlStore::read_manifest(&path).await?;
Ok(QmdlStore {
let manifest = RecordingStore::read_manifest(&path).await?;
Ok(RecordingStore {
path,
manifest,
current_entry: None,
})
}
// Creates a new QmdlStore at the given path. This involves creating a dir
// Creates a new RecordingStore at the given path. This involves creating a dir
// and writing an empty manifest.
pub async fn create<P>(path: P) -> Result<Self, QmdlStoreError> where P: AsRef<Path> {
pub async fn create<P>(path: P) -> Result<Self, RecordingStoreError> where P: AsRef<Path> {
let manifest_path = path.as_ref().join("manifest.toml");
fs::create_dir_all(&path).await
.map_err(QmdlStoreError::OpenDirError)?;
.map_err(RecordingStoreError::OpenDirError)?;
let mut manifest_file = File::create(&manifest_path).await
.map_err(QmdlStoreError::WriteManifestError)?;
.map_err(RecordingStoreError::WriteManifestError)?;
let empty_manifest = Manifest { entries: Vec::new() };
let empty_manifest_contents = toml::to_string_pretty(&empty_manifest)
.expect("failed to serialize manifest");
manifest_file.write_all(empty_manifest_contents.as_bytes()).await
.map_err(QmdlStoreError::WriteManifestError)?;
QmdlStore::load(path).await
.map_err(RecordingStoreError::WriteManifestError)?;
RecordingStore::load(path).await
}
async fn read_manifest<P>(path: P) -> Result<Manifest, QmdlStoreError> where P: AsRef<Path> {
async fn read_manifest<P>(path: P) -> Result<Manifest, RecordingStoreError> where P: AsRef<Path> {
let manifest_path = path.as_ref().join("manifest.toml");
let file_contents = fs::read_to_string(&manifest_path).await
.map_err(QmdlStoreError::ReadManifestError)?;
.map_err(RecordingStoreError::ReadManifestError)?;
toml::from_str(&file_contents)
.map_err(QmdlStoreError::ParseManifestError)
.map_err(RecordingStoreError::ParseManifestError)
}
// Closes the current entry (if needed), creates a new entry based on the
// current time, and updates the manifest
pub async fn new_entry(&mut self) -> Result<File, QmdlStoreError> {
// current time, and updates the manifest. Returns a tuple of the entry's
// newly created QMDL file and analysis file.
pub async fn new_entry(&mut self) -> Result<(File, File), RecordingStoreError> {
// if we've already got an entry open, close it
if self.current_entry.is_some() {
self.close_current_entry().await?;
}
let new_entry = ManifestEntry::new();
let mut file_path = self.path.join(&new_entry.name);
file_path.set_extension("qmdl");
let file = File::options()
let qmdl_filepath = new_entry.get_qmdl_filepath(&self.path);
let qmdl_file = File::options()
.create(true)
.write(true)
.open(&file_path).await
.map_err(QmdlStoreError::CreateFileError)?;
.open(&qmdl_filepath).await
.map_err(RecordingStoreError::CreateFileError)?;
let analysis_filepath = new_entry.get_analysis_filepath(&self.path);
let analysis_file = File::options()
.create(true)
.write(true)
.open(&analysis_filepath).await
.map_err(RecordingStoreError::CreateFileError)?;
self.manifest.entries.push(new_entry);
self.current_entry = Some(self.manifest.entries.len() - 1);
self.write_manifest().await?;
Ok(file)
Ok((qmdl_file, analysis_file))
}
// Returns the corresponding QMDL file for a given entry
pub async fn open_entry(&self, entry: &ManifestEntry) -> Result<File, QmdlStoreError> {
let mut file_path = self.path.join(&entry.name);
file_path.set_extension("qmdl");
File::open(file_path).await
.map_err(QmdlStoreError::ReadFileError)
pub async fn open_entry_qmdl(&self, entry: &ManifestEntry) -> Result<File, RecordingStoreError> {
File::open(entry.get_qmdl_filepath(&self.path)).await
.map_err(RecordingStoreError::ReadFileError)
}
// Returns the corresponding QMDL file for a given entry
pub async fn open_entry_analysis(&self, entry: &ManifestEntry) -> Result<File, RecordingStoreError> {
File::open(entry.get_analysis_filepath(&self.path)).await
.map_err(RecordingStoreError::ReadFileError)
}
// Unsets the current entry
pub async fn close_current_entry(&mut self) -> Result<(), QmdlStoreError> {
pub async fn close_current_entry(&mut self) -> Result<(), RecordingStoreError> {
match self.current_entry {
Some(_) => {
self.current_entry = None;
Ok(())
},
None => Err(QmdlStoreError::NoCurrentEntry)
None => Err(RecordingStoreError::NoCurrentEntry)
}
}
// Sets the given entry's size and updates the last_message_time to now, updating the manifest
pub async fn update_entry(&mut self, entry_index: usize, size_bytes: usize) -> Result<(), QmdlStoreError> {
self.manifest.entries[entry_index].size_bytes = size_bytes;
pub async fn update_entry_qmdl_size(&mut self, entry_index: usize, size_bytes: usize) -> Result<(), RecordingStoreError> {
self.manifest.entries[entry_index].qmdl_size_bytes = size_bytes;
self.manifest.entries[entry_index].last_message_time = Some(Local::now());
self.write_manifest().await
}
async fn write_manifest(&mut self) -> Result<(), QmdlStoreError> {
// Sets the given entry's analysis file size
pub async fn update_entry_analysis_size(&mut self, entry_index: usize, size_bytes: usize) -> Result<(), RecordingStoreError> {
self.manifest.entries[entry_index].analysis_size_bytes = size_bytes;
self.write_manifest().await
}
async fn write_manifest(&mut self) -> Result<(), RecordingStoreError> {
let mut manifest_file = File::options()
.write(true)
.open(self.path.join("manifest.toml")).await
.map_err(QmdlStoreError::WriteManifestError)?;
.map_err(RecordingStoreError::WriteManifestError)?;
let manifest_contents = toml::to_string_pretty(&self.manifest)
.expect("failed to serialize manifest");
manifest_file.write_all(manifest_contents.as_bytes()).await
.map_err(QmdlStoreError::WriteManifestError)?;
.map_err(RecordingStoreError::WriteManifestError)?;
Ok(())
}
@@ -164,6 +194,11 @@ impl QmdlStore {
.find(|entry| entry.name == name)
.cloned()
}
pub fn get_current_entry(&self) -> Option<&ManifestEntry> {
let entry_index = self.current_entry?;
self.manifest.entries.get(entry_index)
}
}
#[cfg(test)]
@@ -174,36 +209,36 @@ mod tests {
#[tokio::test]
async fn test_load_from_empty_dir() {
let dir = TempDir::new("qmdl_store_test").unwrap();
assert!(!QmdlStore::exists(dir.path()).await.unwrap());
let _created_store = QmdlStore::create(dir.path()).await.unwrap();
assert!(QmdlStore::exists(dir.path()).await.unwrap());
let loaded_store = QmdlStore::load(dir.path()).await.unwrap();
assert!(!RecordingStore::exists(dir.path()).await.unwrap());
let _created_store = RecordingStore::create(dir.path()).await.unwrap();
assert!(RecordingStore::exists(dir.path()).await.unwrap());
let loaded_store = RecordingStore::load(dir.path()).await.unwrap();
assert_eq!(loaded_store.manifest.entries.len(), 0);
}
#[tokio::test]
async fn test_creating_updating_and_closing_entries() {
let dir = TempDir::new("qmdl_store_test").unwrap();
let mut store = QmdlStore::create(dir.path()).await.unwrap();
let mut store = RecordingStore::create(dir.path()).await.unwrap();
let _ = store.new_entry().await.unwrap();
let entry_index = store.current_entry.unwrap();
assert_eq!(QmdlStore::read_manifest(dir.path()).await.unwrap(), store.manifest);
assert_eq!(RecordingStore::read_manifest(dir.path()).await.unwrap(), store.manifest);
assert!(store.manifest.entries[entry_index].last_message_time.is_none());
store.update_entry(entry_index, 1000).await.unwrap();
store.update_entry_qmdl_size(entry_index, 1000).await.unwrap();
let entry = store.entry_for_name(&store.manifest.entries[entry_index].name).unwrap();
assert!(entry.last_message_time.is_some());
assert_eq!(store.manifest.entries[entry_index].size_bytes, 1000);
assert_eq!(QmdlStore::read_manifest(dir.path()).await.unwrap(), store.manifest);
assert_eq!(store.manifest.entries[entry_index].qmdl_size_bytes, 1000);
assert_eq!(RecordingStore::read_manifest(dir.path()).await.unwrap(), store.manifest);
store.close_current_entry().await.unwrap();
assert!(matches!(store.close_current_entry().await, Err(QmdlStoreError::NoCurrentEntry)));
assert!(matches!(store.close_current_entry().await, Err(RecordingStoreError::NoCurrentEntry)));
}
#[tokio::test]
async fn test_repeated_new_entries() {
let dir = TempDir::new("qmdl_store_test").unwrap();
let mut store = QmdlStore::create(dir.path()).await.unwrap();
let mut store = RecordingStore::create(dir.path()).await.unwrap();
let _ = store.new_entry().await.unwrap();
let entry_index = store.current_entry.unwrap();
let _ = store.new_entry().await.unwrap();