From 76eb858205964eb186d3f689dd71153d04d54531 Mon Sep 17 00:00:00 2001 From: Caden Date: Wed, 21 May 2025 22:39:14 -0400 Subject: [PATCH] feat(db): rough database implementation --- .env.example | 6 +- Cargo.lock | 107 +++++++++++++ Cargo.toml | 3 +- diesel.toml | 9 ++ migrations/.keep | 0 .../down.sql | 6 + .../up.sql | 36 +++++ .../down.sql | 1 + .../up.sql | 7 + src/api.rs | 36 +++-- src/crypto/symmetric_key_manager.rs | 88 +++++++---- src/db/connection.rs | 13 ++ src/db/mod.rs | 5 + src/db/models.rs | 14 ++ src/db/schema.rs | 12 ++ src/logging.rs | 36 +++-- src/main.rs | 22 ++- src/models.rs | 57 ++++--- src/node.rs | 149 ++++++++++++++---- src/uuid.rs | 13 +- 20 files changed, 505 insertions(+), 115 deletions(-) create mode 100644 diesel.toml create mode 100644 migrations/.keep create mode 100644 migrations/00000000000000_diesel_initial_setup/down.sql create mode 100644 migrations/00000000000000_diesel_initial_setup/up.sql create mode 100644 migrations/2025-05-21-234457_create_encrypted_indicators/down.sql create mode 100644 migrations/2025-05-21-234457_create_encrypted_indicators/up.sql create mode 100644 src/db/connection.rs create mode 100644 src/db/mod.rs create mode 100644 src/db/models.rs create mode 100644 src/db/schema.rs diff --git a/.env.example b/.env.example index c7eec65..a48cb17 100644 --- a/.env.example +++ b/.env.example @@ -4,5 +4,7 @@ DTIM__DEFAULT__PORT=3030 DTIM__DEFAULT__LOG_LEVEL="info" # Sensitive info (.env only) -DTIM__DEFAULT__STORAGE__DATABASE_URL="postgres://user:pass@localhost/db" -DTIM__DEFAULT__WATCHERS__VIRUSTOTAL_API_KEY="your_api_key" \ No newline at end of file +DATABASE_URL="postgres://postgres:@localhost:5432/postgres" # For Diesel CLI + +DTIM__DEFAULT__STORAGE__DATABASE_URL=${DATABASE_URL} +DTIM__DEFAULT__WATCHERS__VIRUSTOTAL_API_KEY="your_api_key" diff --git a/Cargo.lock b/Cargo.lock index 2d89324..409fb19 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -340,6 +340,12 @@ version = "3.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1628fb46dfa0b37568d12e5edd512553eccf6a22a78e8bde00bb4aed84d5bdbf" +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.10.1" @@ -600,6 +606,43 @@ dependencies = [ "serde", ] +[[package]] +name = "diesel" +version = "2.2.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff3e1edb1f37b4953dd5176916347289ed43d7119cc2e6c7c3f7849ff44ea506" +dependencies = [ + "bitflags", + "byteorder", + "chrono", + "diesel_derives", + "itoa", + "pq-sys", + "r2d2", +] + +[[package]] +name = "diesel_derives" +version = "2.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68d4216021b3ea446fd2047f5c8f8fe6e98af34508a254a01e4d6bc1e844f84d" +dependencies = [ + "diesel_table_macro_syntax", + "dsl_auto_type", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "diesel_table_macro_syntax" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "209c735641a413bc68c4923a9d6ad4bcb3ca306b794edaa7eb0b3228a99ffb25" +dependencies = [ + "syn", +] + [[package]] name = "digest" version = "0.10.7" @@ -625,6 +668,20 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +[[package]] +name = "dsl_auto_type" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139ae9aca7527f85f26dd76483eb38533fd84bd571065da1739656ef71c5ff5b" +dependencies = [ + "darling", + "either", + "heck", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "dtim" version = "0.1.0" @@ -636,6 +693,7 @@ dependencies = [ "base64 0.22.1", "chrono", "config", + "diesel", "dotenvy", "ed25519-dalek", "env_logger", @@ -925,6 +983,12 @@ dependencies = [ "hashbrown 0.15.3", ] +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "hex" version = "0.4.3" @@ -1444,6 +1508,16 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" +[[package]] +name = "pq-sys" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41c852911b98f5981956037b2ca976660612e548986c30af075e753107bc3400" +dependencies = [ + "libc", + "vcpkg", +] + [[package]] name = "prettyplease" version = "0.2.32" @@ -1478,6 +1552,17 @@ version = "5.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "74765f6d916ee2faa39bc8e68e4f3ed8949b48cccdac59983d287a7cb71ce9c5" +[[package]] +name = "r2d2" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" +dependencies = [ + "log", + "parking_lot", + "scheduled-thread-pool", +] + [[package]] name = "rand_core" version = "0.6.4" @@ -1653,6 +1738,15 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" +[[package]] +name = "scheduled-thread-pool" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" +dependencies = [ + "parking_lot", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -1759,6 +1853,12 @@ dependencies = [ "syn", ] +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + [[package]] name = "sha2" version = "0.10.9" @@ -2112,8 +2212,15 @@ checksum = "458f7a779bf54acc9f347480ac654f68407d3aab21269a6e3c9f922acd9e2da9" dependencies = [ "getrandom 0.3.3", "serde", + "sha1_smol", ] +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + [[package]] name = "version_check" version = "0.9.5" diff --git a/Cargo.toml b/Cargo.toml index dc1789d..b1eaed1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,7 +8,7 @@ default-run = "dtim" axum = { version = "0.8", features = ["http2", "tokio"] } axum-server = { version = "0.7", features = ["tls-rustls"] } chrono = { version = "0.4", features = ["serde"] } -uuid = { version = "1.16", features = ["v7", "serde"] } +uuid = { version = "1.16", features = ["v5", "v7", "serde"] } serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["preserve_order"] } serde_with = "3.12" @@ -28,3 +28,4 @@ async-trait = "0.1" http-body-util = "0.1" dotenvy = "0.15" once_cell = "1.21" +diesel = { version = "2.2.0", features = ["postgres", "r2d2", "chrono"] } diff --git a/diesel.toml b/diesel.toml new file mode 100644 index 0000000..6623fad --- /dev/null +++ b/diesel.toml @@ -0,0 +1,9 @@ +# For documentation on how to configure this file, +# see https://diesel.rs/guides/configuring-diesel-cli + +[print_schema] +file = "src/db/schema.rs" +custom_type_derives = ["diesel::query_builder::QueryId", "Clone"] + +[migrations_directory] +dir = "/Users/caden/Developer/dtim/migrations" diff --git a/migrations/.keep b/migrations/.keep new file mode 100644 index 0000000..e69de29 diff --git a/migrations/00000000000000_diesel_initial_setup/down.sql b/migrations/00000000000000_diesel_initial_setup/down.sql new file mode 100644 index 0000000..a9f5260 --- /dev/null +++ b/migrations/00000000000000_diesel_initial_setup/down.sql @@ -0,0 +1,6 @@ +-- This file was automatically created by Diesel to setup helper functions +-- and other internal bookkeeping. This file is safe to edit, any future +-- changes will be added to existing projects as new migrations. + +DROP FUNCTION IF EXISTS diesel_manage_updated_at(_tbl regclass); +DROP FUNCTION IF EXISTS diesel_set_updated_at(); diff --git a/migrations/00000000000000_diesel_initial_setup/up.sql b/migrations/00000000000000_diesel_initial_setup/up.sql new file mode 100644 index 0000000..d68895b --- /dev/null +++ b/migrations/00000000000000_diesel_initial_setup/up.sql @@ -0,0 +1,36 @@ +-- This file was automatically created by Diesel to setup helper functions +-- and other internal bookkeeping. This file is safe to edit, any future +-- changes will be added to existing projects as new migrations. + + + + +-- Sets up a trigger for the given table to automatically set a column called +-- `updated_at` whenever the row is modified (unless `updated_at` was included +-- in the modified columns) +-- +-- # Example +-- +-- ```sql +-- CREATE TABLE users (id SERIAL PRIMARY KEY, updated_at TIMESTAMP NOT NULL DEFAULT NOW()); +-- +-- SELECT diesel_manage_updated_at('users'); +-- ``` +CREATE OR REPLACE FUNCTION diesel_manage_updated_at(_tbl regclass) RETURNS VOID AS $$ +BEGIN + EXECUTE format('CREATE TRIGGER set_updated_at BEFORE UPDATE ON %s + FOR EACH ROW EXECUTE PROCEDURE diesel_set_updated_at()', _tbl); +END; +$$ LANGUAGE plpgsql; + +CREATE OR REPLACE FUNCTION diesel_set_updated_at() RETURNS trigger AS $$ +BEGIN + IF ( + NEW IS DISTINCT FROM OLD AND + NEW.updated_at IS NOT DISTINCT FROM OLD.updated_at + ) THEN + NEW.updated_at := current_timestamp; + END IF; + RETURN NEW; +END; +$$ LANGUAGE plpgsql; diff --git a/migrations/2025-05-21-234457_create_encrypted_indicators/down.sql b/migrations/2025-05-21-234457_create_encrypted_indicators/down.sql new file mode 100644 index 0000000..2befc1b --- /dev/null +++ b/migrations/2025-05-21-234457_create_encrypted_indicators/down.sql @@ -0,0 +1 @@ +DROP TABLE encrypted_indicators; diff --git a/migrations/2025-05-21-234457_create_encrypted_indicators/up.sql b/migrations/2025-05-21-234457_create_encrypted_indicators/up.sql new file mode 100644 index 0000000..b5a7574 --- /dev/null +++ b/migrations/2025-05-21-234457_create_encrypted_indicators/up.sql @@ -0,0 +1,7 @@ +CREATE TABLE encrypted_indicators ( + id CHAR(64) PRIMARY KEY, + ciphertext BYTEA NOT NULL, + nonce BYTEA NOT NULL, + mac BYTEA NOT NULL, + tlp_level TEXT NOT NULL +); diff --git a/src/api.rs b/src/api.rs index b2732a6..2a1d0d7 100644 --- a/src/api.rs +++ b/src/api.rs @@ -13,19 +13,17 @@ use ed25519_dalek::VerifyingKey; use http_body_util::BodyExt as _; use rustls::ServerConfig; use serde::Serialize; -use std::{str::FromStr as _, sync::Arc}; +use std::sync::Arc; use tokio::sync::Mutex; +use crate::models::ThreatIndicator; use crate::{ crypto::MeshIdentity, + db::models::EncryptedIndicator, error::{ApiError, ApiErrorResponse}, node::{Node, NodePeer}, }; use crate::{crypto::SymmetricKeyManager, models::StixBundle}; -use crate::{ - models::{EncryptedThreatIndicator, ThreatIndicator}, - uuid::Uuid, -}; #[derive(Clone)] pub struct AppState { @@ -200,13 +198,14 @@ struct GossipIndicatorsResponse { async fn gossip_indicators_handler( State(state): State>, - Json(indicators): Json>, + Json(indicators): Json>, ) -> ApiResponse { let mut node = state.node.lock().await; let mut count = 0; for encrypted in indicators { if let Ok(indicator) = ThreatIndicator::decrypt(&encrypted, &state.key_mgr) { - node.add_or_increment_indicator(indicator); + node.add_or_increment_indicator(indicator) + .map_err(|_| ApiError::INTERNAL_SERVER_ERROR)?; count += 1; } } @@ -229,7 +228,9 @@ async fn get_public_indicators_handler( State(state): State>, ) -> ApiResponse { let node = state.node.lock().await; - let indicators = node.list_indicators_by_tlp(crate::models::TlpLevel::White); + let indicators = node + .list_indicators_by_tlp(crate::models::TlpLevel::White) + .map_err(|_| ApiError::INTERNAL_SERVER_ERROR)?; // Return anonymized (open) view Ok(( StatusCode::OK, @@ -244,7 +245,9 @@ async fn get_private_indicators_handler( State(state): State>, ) -> ApiResponse { let node = state.node.lock().await; - let indicators = node.list_indicators_by_tlp(crate::models::TlpLevel::Red); + let indicators = node + .list_indicators_by_tlp(crate::models::TlpLevel::Red) + .map_err(|_| ApiError::INTERNAL_SERVER_ERROR)?; // TODO: Ensure the node is an authenticated recipient for private indicators // Return anonymized (moderate) view Ok(( @@ -261,8 +264,9 @@ async fn get_indicator_by_id_handler( Path(id): Path, ) -> ApiResponse { let node = state.node.lock().await; - let id = Uuid::from_str(&id).map_err(|_| ApiError::INVALID_INDICATOR_ID)?; - let indicator = node.get_indicator_by_id(&id).ok_or(ApiError::NOT_FOUND)?; + let indicator = node + .get_indicator_by_id(&id) + .map_err(|_| ApiError::NOT_FOUND)?; Ok((StatusCode::OK, Json(indicator.to_json(node.get_level())))) } @@ -319,7 +323,12 @@ async fn taxii_get_objects_handler( Path(_collection_id): Path, ) -> ApiResponse { let node = state.node.lock().await; - let stix_objects = node.list_objects_by_tlp(crate::models::TlpLevel::White); + let stix_objects = node + .list_objects_by_tlp(crate::models::TlpLevel::White) + .map_err(|error| { + println!("Error: {:?}", error); + ApiError::INTERNAL_SERVER_ERROR + })?; let bundle = StixBundle::new(stix_objects); Ok((StatusCode::OK, Json(bundle.to_stix()))) } @@ -331,6 +340,7 @@ async fn taxii_post_objects_handler( ) -> ApiResponse { let indicator = ThreatIndicator::from_stix(stix).map_err(|_| ApiError::INVALID_STIX_OBJECT)?; let mut node = state.node.lock().await; - node.add_indicator(indicator); + node.add_indicator(indicator) + .map_err(|_| ApiError::INTERNAL_SERVER_ERROR)?; Ok((StatusCode::OK, Json(serde_json::json!({ "status": "ok" })))) } diff --git a/src/crypto/symmetric_key_manager.rs b/src/crypto/symmetric_key_manager.rs index 4265314..18ab889 100644 --- a/src/crypto/symmetric_key_manager.rs +++ b/src/crypto/symmetric_key_manager.rs @@ -1,9 +1,16 @@ use aes_gcm::aead::rand_core::RngCore; use aes_gcm::aead::{Aead, KeyInit, OsRng}; use aes_gcm::{Aes256Gcm, Key, Nonce}; -use base64::prelude::*; +use std::fs::{self, File}; +use std::io::{Read, Write}; +use std::path::Path; use std::time::{SystemTime, UNIX_EPOCH}; +const SYMM_KEY_PATH: &str = "data/keys/symm.key"; +const SYMM_PREV_KEY_PATH: &str = "data/keys/symm_prev.key"; + +pub type EncryptedData = (Vec, Vec, Vec); + #[derive(Clone, Debug)] pub struct SymmetricKeyManager { current_key: Key, @@ -13,18 +20,40 @@ pub struct SymmetricKeyManager { } impl SymmetricKeyManager { - pub fn new(rotation_days: u64) -> Self { - let current_key = Self::generate_key(); + pub fn load_or_generate(rotation_days: u64) -> std::io::Result { + fs::create_dir_all("data/keys")?; + + let mut current_bytes = [0u8; 32]; + let mut prev_bytes = [0u8; 32]; + + let current_key = if Path::new(SYMM_KEY_PATH).exists() { + File::open(SYMM_KEY_PATH)?.read_exact(&mut current_bytes)?; + Key::::from_slice(¤t_bytes).to_owned() + } else { + let key = Self::generate_key(); + let mut file = File::create(SYMM_KEY_PATH)?; + file.write_all(key.as_slice())?; + key + }; + + let previous_key = if Path::new(SYMM_PREV_KEY_PATH).exists() { + File::open(SYMM_PREV_KEY_PATH)?.read_exact(&mut prev_bytes)?; + Some(Key::::from_slice(&prev_bytes).to_owned()) + } else { + None + }; + let key_rotation_time = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(); - Self { + + Ok(Self { current_key, - previous_key: None, + previous_key, key_rotation_time, rotation_interval: rotation_days * 24 * 60 * 60, - } + }) } fn generate_key() -> Key { @@ -33,7 +62,21 @@ impl SymmetricKeyManager { Key::::from_slice(&key_bytes).to_owned() } - pub fn rotate_key(&mut self) { + pub fn save_keys(&self) -> std::io::Result<()> { + fs::create_dir_all("data/keys")?; + let mut file = File::create(SYMM_KEY_PATH)?; + file.write_all(self.current_key.as_slice())?; + + if let Some(prev) = &self.previous_key { + let mut prev_file = File::create(SYMM_PREV_KEY_PATH)?; + prev_file.write_all(prev.as_slice())?; + } else if Path::new(SYMM_PREV_KEY_PATH).exists() { + fs::remove_file(SYMM_PREV_KEY_PATH)?; + } + Ok(()) + } + + pub fn rotate_key(&mut self) -> std::io::Result<()> { let now = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() @@ -42,11 +85,13 @@ impl SymmetricKeyManager { self.previous_key = Some(self.current_key); self.current_key = Self::generate_key(); self.key_rotation_time = now; + self.save_keys()?; } + Ok(()) } - pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<(String, String, String), String> { - self.rotate_key(); + pub fn encrypt(&mut self, plaintext: &[u8]) -> Result { + self.rotate_key()?; let cipher = Aes256Gcm::new(&self.current_key); let mut nonce_bytes = [0u8; 12]; OsRng.fill_bytes(&mut nonce_bytes); @@ -54,33 +99,20 @@ impl SymmetricKeyManager { match cipher.encrypt(nonce, plaintext) { Ok(mut ciphertext_and_tag) => { let tag = ciphertext_and_tag.split_off(ciphertext_and_tag.len() - 16); - Ok(( - BASE64_STANDARD.encode(ciphertext_and_tag), - BASE64_STANDARD.encode(nonce_bytes), - BASE64_STANDARD.encode(tag), - )) + Ok((ciphertext_and_tag, nonce_bytes.to_vec(), tag)) } - Err(_) => Err("Encryption failure".to_string()), + Err(_) => Err(std::io::Error::other("Encryption failure")), } } pub fn decrypt( &self, - ciphertext_b64: &str, - nonce_b64: &str, - tag_b64: &str, + mut ciphertext: Vec, + nonce: Vec, + tag: Vec, ) -> Result, String> { - let mut ciphertext = BASE64_STANDARD - .decode(ciphertext_b64) - .map_err(|_| "Invalid base64 ciphertext")?; - let nonce_bytes = BASE64_STANDARD - .decode(nonce_b64) - .map_err(|_| "Invalid base64 nonce")?; - let tag = BASE64_STANDARD - .decode(tag_b64) - .map_err(|_| "Invalid base64 tag")?; ciphertext.extend_from_slice(&tag); - let nonce = Nonce::from_slice(&nonce_bytes); + let nonce = Nonce::from_slice(&nonce); let try_decrypt = |key: &Key| { let cipher = Aes256Gcm::new(key); cipher diff --git a/src/db/connection.rs b/src/db/connection.rs new file mode 100644 index 0000000..003c5cc --- /dev/null +++ b/src/db/connection.rs @@ -0,0 +1,13 @@ +use diesel::{ + pg::PgConnection, + r2d2::{ConnectionManager, Pool}, +}; +use std::error::Error; + +pub fn get_connection_pool( + conn_str: &str, +) -> Result>, Box> { + let manager = ConnectionManager::::new(conn_str); + let pool = Pool::builder().build(manager)?; + Ok(pool) +} diff --git a/src/db/mod.rs b/src/db/mod.rs new file mode 100644 index 0000000..e57d22b --- /dev/null +++ b/src/db/mod.rs @@ -0,0 +1,5 @@ +pub mod connection; +pub mod models; +pub mod schema; + +pub use connection::get_connection_pool; diff --git a/src/db/models.rs b/src/db/models.rs new file mode 100644 index 0000000..cbf6388 --- /dev/null +++ b/src/db/models.rs @@ -0,0 +1,14 @@ +use crate::db::schema::encrypted_indicators; +use diesel::prelude::*; +use diesel::Selectable; +use serde::{Deserialize, Serialize}; + +#[derive(Queryable, Insertable, Selectable, Serialize, Deserialize, Debug)] +#[diesel(table_name = encrypted_indicators)] +pub struct EncryptedIndicator { + pub id: String, + pub ciphertext: Vec, + pub nonce: Vec, + pub mac: Vec, + pub tlp_level: String, +} diff --git a/src/db/schema.rs b/src/db/schema.rs new file mode 100644 index 0000000..5df65f8 --- /dev/null +++ b/src/db/schema.rs @@ -0,0 +1,12 @@ +// @generated automatically by Diesel CLI. + +diesel::table! { + encrypted_indicators (id) { + #[max_length = 64] + id -> Bpchar, + ciphertext -> Bytea, + nonce -> Bytea, + mac -> Bytea, + tlp_level -> Text, + } +} diff --git a/src/logging.rs b/src/logging.rs index 9a46c47..9b48692 100644 --- a/src/logging.rs +++ b/src/logging.rs @@ -1,3 +1,5 @@ +use base64::prelude::BASE64_STANDARD; +use base64::Engine as _; use chrono::Utc; use log::{error, Level, LevelFilter, Metadata, Record}; use std::fs::{self, OpenOptions}; @@ -37,7 +39,12 @@ impl EncryptedLogger { .key_mgr .encrypt(log_entry.as_bytes()) .map_err(|e| std::io::Error::other(format!("Encryption failed: {}", e)))?; - let encrypted_entry = format!("{}\n{}\n{}\n", ciphertext, nonce, mac); + let encrypted_entry = format!( + "{}\n{}\n{}\n", + BASE64_STANDARD.encode(ciphertext), + BASE64_STANDARD.encode(nonce), + BASE64_STANDARD.encode(mac) + ); let filename = format!("{}.log", Utc::now().format("%Y-%m-%d")); let log_file = self.log_path.join(filename); @@ -66,15 +73,24 @@ impl EncryptedLogger { let mut lines = content.lines().peekable(); while lines.peek().is_some() { - let ciphertext = lines.next().ok_or_else(|| { - std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid log format") - })?; - let nonce = lines.next().ok_or_else(|| { - std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid log format") - })?; - let mac = lines.next().ok_or_else(|| { - std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid log format") - })?; + let ciphertext = lines + .next() + .and_then(|line| BASE64_STANDARD.decode(line).ok()) + .ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid log format") + })?; + let nonce = lines + .next() + .and_then(|line| BASE64_STANDARD.decode(line).ok()) + .ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid log format") + })?; + let mac = lines + .next() + .and_then(|line| BASE64_STANDARD.decode(line).ok()) + .ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid log format") + })?; match self.key_mgr.decrypt(ciphertext, nonce, mac) { Ok(decrypted) => { diff --git a/src/main.rs b/src/main.rs index 34f317d..a085b99 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,6 @@ mod api; mod crypto; +mod db; mod error; mod logging; mod models; @@ -76,7 +77,8 @@ fn make_server_config(settings: &settings::Settings) -> Arc Result<(), Box> { init_logging(); let settings = Settings::new()?; - let mut key_mgr = crypto::SymmetricKeyManager::new(settings.tls.key_rotation_days); + let mut key_mgr = + crypto::SymmetricKeyManager::load_or_generate(settings.tls.key_rotation_days)?; let tls_config = make_server_config(&settings); let logger = logging::EncryptedLogger::new( @@ -91,7 +93,9 @@ async fn main() -> Result<(), Box> { }), )?; - let node = node::Node::new(logger, settings.privacy)?; + println!("Database URL: {:?}", settings.storage.database_url); + let db_pool = db::get_connection_pool(&settings.storage.database_url)?; + let node = node::Node::new(db_pool, key_mgr.clone(), logger, settings.privacy)?; let id = node.get_id(); println!("Node ID: {:?}", id); @@ -130,19 +134,25 @@ async fn main() -> Result<(), Box> { vec!["malicious-activity".to_string()], models::TlpLevel::White, None, + ) + .unwrap(); + + let encrypted = indicator.encrypt(&mut key_mgr).unwrap(); + println!( + "Encrypted: {:?}", + serde_json::to_string(&encrypted).unwrap() ); - let encrypted = indicator.encrypt(&mut key_mgr); - println!("Encrypted: {:?}", encrypted); + let _ = key_mgr.rotate_key(); - let decrypted = ThreatIndicator::decrypt(&encrypted.unwrap(), &key_mgr); + let decrypted = ThreatIndicator::decrypt(&encrypted, &key_mgr).unwrap(); println!("Decrypted: {:?}", decrypted); let node = Arc::new(Mutex::new(node)); { let mut node = node.lock().await; - node.add_indicator(indicator.clone()); + let _ = node.add_or_increment_indicator(indicator.clone()); node.bootstrap_peers(settings.network.init_peers.clone()); } diff --git a/src/models.rs b/src/models.rs index 553db92..ab3c4ba 100644 --- a/src/models.rs +++ b/src/models.rs @@ -4,11 +4,12 @@ use std::{ net::IpAddr, }; -use crate::{crypto::SymmetricKeyManager, uuid::Uuid}; +use crate::{crypto::SymmetricKeyManager, db::models::EncryptedIndicator, uuid::Uuid}; use chrono::{DateTime, Utc}; use regex::Regex; use serde::{Deserialize, Serialize}; use serde_json::json; +use sha2::{Digest, Sha256}; use std::hash::{Hash, Hasher}; #[derive(Serialize, Deserialize, Clone, Debug)] @@ -27,13 +28,6 @@ pub struct ThreatIndicator { pub marking_definitions: Vec, } -#[derive(Serialize, Deserialize, Clone, Debug)] -pub struct EncryptedThreatIndicator { - ciphertext: String, - nonce: String, - mac: String, -} - impl ThreatIndicator { pub fn new( indicator_type: IndicatorType, @@ -42,12 +36,14 @@ impl ThreatIndicator { tags: Vec, tlp: TlpLevel, custom_fields: Option>, - ) -> Self { + ) -> Result { + let id = Self::compute_id(&indicator_type, &value, &tlp, &tags) + .map_err(|e| std::io::Error::other(format!("Failed to compute ID: {}", e)))?; let now = Utc::now(); let marking_definition = MarkingDefinition::new("tlp".to_string(), tlp.to_string()); - ThreatIndicator { - id: Uuid::new_v7_from_datetime(now), + Ok(ThreatIndicator { + id, indicator_type, value, confidence, @@ -59,35 +55,58 @@ impl ThreatIndicator { recipients: None, // TODO: Handle recipients for TLP:RED custom_fields, marking_definitions: vec![marking_definition], - } + }) + } + + pub fn get_id(&self) -> String { + self.id.to_string() } - pub fn get_id(&self) -> Uuid { - self.id + pub fn compute_id( + indicator_type: &IndicatorType, + value: &str, + tlp: &TlpLevel, + tags: &[String], + ) -> Result { + let mut hasher = Sha256::new(); + hasher.update(indicator_type.to_string().as_bytes()); + hasher.update(value.as_bytes()); + hasher.update(tlp.to_string().as_bytes()); + for tag in tags { + hasher.update(tag.as_bytes()); + } + let hash = hasher.finalize(); + Ok(Uuid::new_v5_from_hash(&hash)) } pub fn encrypt( &self, key_mgr: &mut SymmetricKeyManager, - ) -> Result { + ) -> Result { let serialized = serde_json::to_vec(self).expect("Failed to serialize ThreatIndicator"); let (ciphertext, nonce, mac) = key_mgr .encrypt(&serialized) .map_err(|e| std::io::Error::other(format!("Encryption failed: {}", e)))?; - Ok(EncryptedThreatIndicator { + Ok(EncryptedIndicator { + id: self.id.to_string(), ciphertext, nonce, mac, + tlp_level: self.tlp.to_string(), }) } pub fn decrypt( - encrypted: &EncryptedThreatIndicator, + encrypted: &EncryptedIndicator, key_mgr: &SymmetricKeyManager, ) -> Result { let decrypted = key_mgr - .decrypt(&encrypted.ciphertext, &encrypted.nonce, &encrypted.mac) + .decrypt( + encrypted.ciphertext.clone(), + encrypted.nonce.clone(), + encrypted.mac.clone(), + ) .map_err(|e| format!("Decryption failed: {}", e))?; serde_json::from_slice(&decrypted) @@ -264,7 +283,7 @@ impl MarkingDefinition { let now = Utc::now(); MarkingDefinition { - id: Uuid::new_v7_from_datetime(now), + id: Uuid::now_v7(), created_at: now, definition_type: definition_type.clone(), definition: json!({ diff --git a/src/node.rs b/src/node.rs index c3050c1..8965507 100644 --- a/src/node.rs +++ b/src/node.rs @@ -1,31 +1,46 @@ use crate::{ - crypto::{self, MeshIdentity}, + crypto::{self, MeshIdentity, SymmetricKeyManager}, + db::{self, models::EncryptedIndicator}, logging::EncryptedLogger, models::{PrivacyLevel, ThreatIndicator, TlpLevel}, settings::PrivacyConfig, - uuid::Uuid, }; use chrono::Utc; +use diesel::{ + pg::PgConnection, + query_dsl::methods::{FilterDsl, FindDsl}, + r2d2::Pool, + OptionalExtension, +}; +use diesel::{r2d2::ConnectionManager, ExpressionMethods}; +use diesel::{RunQueryDsl, SelectableHelper}; use serde::{Deserialize, Serialize}; use std::{collections::HashMap, io}; #[derive(Clone, Debug)] pub struct Node { identity: MeshIdentity, - indicators: HashMap, peers: HashMap, + db_pool: Pool>, + key_mgr: SymmetricKeyManager, logger: EncryptedLogger, privacy_level: PrivacyLevel, allow_custom_fields: bool, } impl Node { - pub fn new(logger: EncryptedLogger, privacy: PrivacyConfig) -> Result { + pub fn new( + db_pool: Pool>, + key_mgr: SymmetricKeyManager, + logger: EncryptedLogger, + privacy: PrivacyConfig, + ) -> Result { let identity = crypto::MeshIdentity::load_or_generate()?; Ok(Node { identity, - indicators: HashMap::new(), peers: HashMap::new(), + db_pool, + key_mgr, logger, privacy_level: match privacy.level.as_str() { "strict" => PrivacyLevel::Strict, @@ -62,49 +77,119 @@ impl Node { ); } - pub fn add_indicator(&mut self, indicator: ThreatIndicator) -> Uuid { - let id = indicator.get_id(); - self.indicators.insert(id, indicator.clone()); - let _ = self.logger.write_log( - log::Level::Info, - &format!("Adding indicator: {:?}", indicator), - ); - id + pub fn add_indicator( + &mut self, + indicator: ThreatIndicator, + ) -> Result> { + use self::db::schema::encrypted_indicators; + + let encrypted_indicator = indicator.encrypt(&mut self.key_mgr)?; + + let mut conn = self.db_pool.get()?; + let res = diesel::insert_into(encrypted_indicators::table) + .values(&encrypted_indicator) + .returning(EncryptedIndicator::as_returning()) + .get_result(&mut conn)?; + + let _ = self + .logger + .write_log(log::Level::Info, &format!("Added indicator: {:?}", res)); + Ok(res) } - pub fn add_or_increment_indicator(&mut self, new_indicator: ThreatIndicator) -> Uuid { - let id = new_indicator.get_id(); - if let Some(existing) = self.indicators.get_mut(&id) { - existing.sightings += 1; - existing.updated_at = Utc::now(); + pub fn add_or_increment_indicator( + &mut self, + new_indicator: ThreatIndicator, + ) -> Result> { + use self::db::schema::encrypted_indicators::dsl::*; + + let indicator_id = new_indicator.get_id(); + + let mut conn = self.db_pool.get()?; + let existing: Option = encrypted_indicators + .find(&indicator_id) + .first::(&mut conn) + .optional()?; + + if let Some(encrypted) = existing { + let mut indicator = ThreatIndicator::decrypt(&encrypted, &self.key_mgr)?; + indicator.sightings += 1; + indicator.updated_at = Utc::now(); + + let new_encrypted = indicator.encrypt(&mut self.key_mgr)?; + let res = diesel::update(encrypted_indicators.find(&indicator_id)) + .set(( + ciphertext.eq(new_encrypted.ciphertext), + nonce.eq(new_encrypted.nonce), + mac.eq(new_encrypted.mac), + )) + .returning(EncryptedIndicator::as_returning()) + .get_result(&mut conn)?; + let _ = self.logger.write_log( log::Level::Info, - &format!("Incrementing indicator: {:?}", existing), + &format!("Incrementing indicator: {:?}", res), ); + Ok(res) } else { - self.add_indicator(new_indicator); + let encrypted = new_indicator.encrypt(&mut self.key_mgr)?; + let res = diesel::insert_into(encrypted_indicators) + .values(&encrypted) + .returning(EncryptedIndicator::as_returning()) + .get_result(&mut conn)?; + + let _ = self + .logger + .write_log(log::Level::Info, &format!("Adding indicator: {:?}", res)); + Ok(res) } - id } pub fn get_level(&self) -> PrivacyLevel { self.privacy_level } - pub fn get_indicator_by_id(&self, id: &Uuid) -> Option<&ThreatIndicator> { - self.indicators.get(id) + pub fn get_indicator_by_id( + &self, + indicator_id: &String, + ) -> Result> { + use self::db::schema::encrypted_indicators::dsl::*; + + let mut conn = self.db_pool.get()?; + let indicator = encrypted_indicators + .find(indicator_id) + .first::(&mut conn)?; + + let indicator = ThreatIndicator::decrypt(&indicator, &self.key_mgr)?; + Ok(indicator) } - pub fn list_indicators_by_tlp(&self, tlp: TlpLevel) -> Vec { - self.indicators - .values() - .filter(|i| i.tlp == tlp) - .cloned() - .collect() + pub fn list_indicators_by_tlp( + &self, + tlp: TlpLevel, + ) -> Result, Box> { + use self::db::schema::encrypted_indicators::dsl::*; + + let mut conn = self.db_pool.get()?; + let indicators = encrypted_indicators + .filter(tlp_level.eq(tlp.to_string())) + .load::(&mut conn)?; + + let indicators = indicators + .iter() + .map(|i| ThreatIndicator::decrypt(i, &self.key_mgr)) + .collect::, std::string::String>>()?; + + Ok(indicators) } - pub fn list_objects_by_tlp(&self, tlp: TlpLevel) -> Vec { - let indicators: Vec<_> = self.list_indicators_by_tlp(tlp); + pub fn list_objects_by_tlp( + &self, + tlp: TlpLevel, + ) -> Result, Box> { + let indicators = self.list_indicators_by_tlp(tlp)?; + + println!("Indicators: {:?}", indicators); let mut stix_indicators: Vec = indicators .iter() @@ -118,7 +203,7 @@ impl Node { stix_indicators.extend(stix_mds); stix_indicators.sort_by(|a, b| a["id"].as_str().unwrap().cmp(b["id"].as_str().unwrap())); - stix_indicators + Ok(stix_indicators) } pub fn get_peers(&self) -> &HashMap { diff --git a/src/uuid.rs b/src/uuid.rs index cc50089..666447e 100644 --- a/src/uuid.rs +++ b/src/uuid.rs @@ -60,12 +60,17 @@ impl Deref for Uuid { impl Uuid { /// Generate a new UUID - pub fn new() -> Self { - Self(uuid::Uuid::now_v7()) + pub fn new_v7() -> Self { + Self::now_v7() } /// Generate a new V7 UUID - pub fn new_v7() -> Self { - Self(uuid::Uuid::now_v7()) + pub fn now_v7() -> Self { + let ctx = TIMESTAMP_CONTEXT.lock().unwrap(); + Self(uuid::Uuid::new_v7(uuid::Timestamp::now(&*ctx))) + } + /// Generate a new V5 UUID + pub fn new_v5_from_hash(hash: &[u8]) -> Self { + Self(uuid::Uuid::new_v5(&uuid::Uuid::NAMESPACE_OID, hash)) } /// Generate a new V7 UUID pub fn new_v7_from_datetime(timestamp: DateTime) -> Self {