milestone cea-sec#1

Author: bshifter

Cargo.lock

@@ -23,6 +23,7 @@ chrono = { version  = "0.4.26", default-features = false, features = ["clock"] }
 encoding_rs = "0.8.32"
 deadpool-postgres = "0.14.0"
 deadpool-sqlite = "0.5.0"
+deadpool-redis = "0.18.0"
 openssl = "0.10.66"
 postgres-openssl = "0.5.0"
 strum = { version = "0.26.1", features = ["derive"] }

common/Cargo.toml

@@ -7,6 +7,7 @@ use crate::{
     bookmark::BookmarkData,
     database::postgres::PostgresDatabase,
     database::sqlite::SQLiteDatabase,
+    database::redis::RedisDatabase,
     heartbeat::{HeartbeatData, HeartbeatsCache},
     settings::Settings,
     subscription::{
@@ -21,6 +22,7 @@ use self::schema::{Migration, Version};
 pub mod postgres;
 pub mod schema;
 pub mod sqlite;
+pub mod redis;
 
 pub type Db = Arc<dyn Database + Send + Sync>;
 
@@ -40,6 +42,12 @@ pub async fn db_from_settings(settings: &Settings) -> Result<Db> {
             schema::postgres::register_migrations(&mut db);
             Ok(Arc::new(db))
         }
+        crate::settings::Database::Redis(redis) => {
+            let db = RedisDatabase::new(redis.connection_string())
+            .await
+            .context("Failed to initialize Redis client")?;
+            Ok(Arc::new(db))
+        }
     }
 }
 

common/src/database/mod.rs

@@ -0,0 +1,278 @@
+// Some of the following code is inspired from
+// https://github.com/SkylerLipthay/schemamama_postgres. As stated by its
+// license (MIT), we include below its copyright notice and permission notice:
+//
+//       The MIT License (MIT)
+//
+//       Copyright (c) 2015 Skyler Lipthay
+//
+//       Permission is hereby granted, free of charge, to any person obtaining a copy
+//       of this software and associated documentation files (the "Software"), to deal
+//       in the Software without restriction, including without limitation the rights
+//       to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+//       copies of the Software, and to permit persons to whom the Software is
+//       furnished to do so, subject to the following conditions:
+//
+//       The above copyright notice and this permission notice shall be included in all
+//       copies or substantial portions of the Software.
+//
+//       THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+//       IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+//       FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+//       AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+//       LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+//       OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+//       SOFTWARE.
+//
+//
+#![allow(unused_imports)]
+use anyhow::{anyhow, ensure, Context, Error, Result};
+use async_trait::async_trait;
+use deadpool_redis::redis::AsyncCommands;
+// use deadpool_sqlite::{Config, Pool, Runtime};
+use deadpool_redis::{Config, Pool, Runtime};
+use log::warn;
+use serde::de::value;
+// use rusqlite::{named_params, params, Connection, OptionalExtension, Row};
+use uuid::Uuid;
+use std::borrow::Borrow;
+use std::collections::btree_map::Entry::Vacant;
+use std::collections::{BTreeMap, BTreeSet};
+use std::str::FromStr;
+use std::sync::Arc;
+use std::time::SystemTime;
+
+use crate::bookmark::BookmarkData;
+use crate::database::Database;
+use crate::heartbeat::{HeartbeatData, HeartbeatsCache};
+use crate::subscription::{
+    ContentFormat, InternalVersion, PrincsFilter, SubscriptionData, SubscriptionMachine, SubscriptionMachineState, SubscriptionStatsCounters, SubscriptionUuid
+};
+
+use super::schema::{Migration, MigrationBase, Version};
+
+pub enum RedisKey {
+    Subscriptions,
+    Machines,
+    Heartbeats,
+    BookMarks
+}
+
+impl RedisKey {
+    pub fn as_str(&self) -> &str {
+        match self {
+            RedisKey::Subscriptions => "subscriptions",
+            RedisKey::Machines => "machines",
+            RedisKey::Heartbeats => "heartbeats",
+            RedisKey::BookMarks => "bookmarks",
+        }
+    }
+}
+
+
+#[allow(unused)]
+pub struct RedisDatabase {
+    pool: Pool,
+}
+
+impl RedisDatabase {
+    pub async fn new(connection_string: &str) -> Result<RedisDatabase> {
+        let config = Config::from_url(connection_string);
+        let pool = config.create_pool(Some(Runtime::Tokio1))?;
+        let db = RedisDatabase {
+            pool,
+        };
+
+        Ok(db)
+    }
+}
+
+#[allow(unused)]
+#[async_trait]
+impl Database for RedisDatabase {
+    async fn get_bookmark(&self, machine: &str, subscription: &str) -> Result<Option<String>> {
+        let mut conn = self.pool.get().await.context("Failed to get Redis connection")?;
+        let key = format!("{}:{}", RedisKey::BookMarks.as_str(), subscription);
+        let machine_s = machine.to_string();
+        let res : Option<String> = conn.hget(key, machine_s).await.context("Failed to get bookmark data")?;
+        Ok(res)
+    }
+
+    async fn get_bookmarks(&self, subscription: &str) -> Result<Vec<BookmarkData>> {
+        todo!()
+    }
+
+    async fn store_bookmark(
+        &self,
+        machine: &str,
+        subscription: &str,
+        bookmark: &str,
+    ) -> Result<()> {
+        let machine_s = machine.to_string();
+        let bookmark_s = bookmark.to_string();
+        let mut conn = self.pool.get().await.context("Failed to get Redis connection")?;
+        let key = format!("{}:{}", RedisKey::BookMarks.as_str(), subscription);
+        let _ : String = conn.hset(key, machine_s, bookmark_s).await.context("Failed to store bookmark data")?;
+        Ok(())
+    }
+
+    async fn delete_bookmarks(
+        &self,
+        machine: Option<&str>,
+        subscription: Option<&str>,
+    ) -> Result<()> {
+        todo!()
+    }
+
+    async fn get_heartbeats_by_machine(
+        &self,
+        machine: &str,
+        subscription: Option<&str>,
+    ) -> Result<Vec<HeartbeatData>> {
+        todo!()
+    }
+
+    async fn get_heartbeats_by_ip(
+        &self,
+        ip: &str,
+        subscription: Option<&str>,
+    ) -> Result<Vec<HeartbeatData>> {
+        todo!()
+    }
+
+    async fn get_heartbeats(&self) -> Result<Vec<HeartbeatData>> {
+        todo!()
+    }
+
+    async fn get_heartbeats_by_subscription(
+        &self,
+        subscription: &str,
+    ) -> Result<Vec<HeartbeatData>> {
+        todo!()
+    }
+
+    async fn store_heartbeat(
+        &self,
+        machine: &str,
+        ip: String,
+        subscription: &str,
+        is_event: bool,
+    ) -> Result<()> {
+        todo!()
+    }
+
+    async fn store_heartbeats(&self, heartbeats: &HeartbeatsCache) -> Result<()> {
+        todo!()
+    }
+
+    async fn get_subscriptions(&self) -> Result<Vec<SubscriptionData>> {
+        let mut conn = self.pool.get().await.context("Failed to get Redis connection")?;
+
+        let keys: Vec<String> = conn.hkeys(RedisKey::Subscriptions.as_str()).await.context("Failed to get subscription keys")?;
+
+        let mut subscriptions = Vec::new();
+
+        for key in keys {
+            let subscription_json: Option<String> = conn.hget(RedisKey::Subscriptions.as_str(), &key).await.context("Failed to get subscription data")?;
+
+            if let Some(subscription_json) = subscription_json {
+                match serde_json::from_str::<SubscriptionData>(&subscription_json) {
+                    Ok(subscription) => subscriptions.push(subscription),
+                    Err(err) => {
+                        log::warn!("Failed to deserialize subscription data for key {}: {}", key, err);
+                    }
+                }
+            } else {
+                log::warn!("No subscription found for key: {}", key);
+            }
+        }
+
+        Ok(subscriptions)
+    }
+
+    async fn get_subscription_by_identifier(
+        &self,
+        identifier: &str,
+    ) -> Result<Option<SubscriptionData>> {
+        let mut conn = self.pool.get().await.context("Failed to get Redis connection")?;
+        let key = identifier;
+        let result: Option<String> = conn.hget(RedisKey::Subscriptions.as_str(), key).await.context("Failed to get subscription data")?;
+        if result.is_none() {
+            return Ok(None);
+        }
+        let subscription: SubscriptionData = serde_json::from_str(&result.unwrap()).context("Failed to deserialize subscription data")?;
+        Ok(Some(subscription))
+    }
+
+    async fn store_subscription(&self, subscription: &SubscriptionData) -> Result<()> {
+        let mut conn = self.pool.get().await.context("Failed to get Redis connection")?;
+        let key = format!("{}", subscription.uuid());
+        // let key = format!("{}", subscription.name());
+        let subs_json = serde_json::to_string(subscription).context("Failed to serialize subscription data")?;
+        let _ : String = conn.hset(RedisKey::Subscriptions.as_str(), key, &subs_json).await.context("Failed to store subscription data")?;
+        Ok(())
+    }
+
+    async fn delete_subscription(&self, uuid: &str) -> Result<()> {
+        let mut conn = self.pool.get().await.context("Failed to get Redis connection")?;
+
+        let key = uuid;
+
+        let result: usize = conn.hdel(RedisKey::Subscriptions.as_str(), key)
+            .await
+            .context("Failed to delete subscription data")?;
+
+        if result == 0 {
+            log::warn!("No subscription found for UUID: {}", uuid);
+        }
+
+        Ok(())
+    }
+
+    /// Fails if `setup_schema` hasn't previously been called or if the query otherwise fails.
+    async fn current_version(&self) -> Result<Option<Version>> {
+        todo!()
+    }
+
+    /// Fails if `setup_schema` hasn't previously been called or if the query otherwise fails.
+    async fn migrated_versions(&self) -> Result<BTreeSet<Version>> {
+        Ok(BTreeSet::<i64>::new())
+    }
+
+    /// Fails if `setup_schema` hasn't previously been called or if the migration otherwise fails.
+    async fn apply_migration(&self, version: Version) -> Result<()> {
+        todo!()
+    }
+
+    /// Fails if `setup_schema` hasn't previously been called or if the migration otherwise fails.
+    async fn revert_migration(&self, version: Version) -> Result<()> {
+        todo!()
+    }
+
+    /// Create the tables required to keep track of schema state. If the tables already
+    /// exist, this function has no operation.
+    async fn setup_schema(&self) -> Result<()> {
+        todo!()
+    }
+
+    async fn migrations(&self) -> BTreeMap<Version, Arc<dyn Migration + Send + Sync>> {
+        BTreeMap::new()
+    }
+
+    async fn get_stats(
+        &self,
+        subscription: &str,
+        start_time: i64,
+    ) -> Result<SubscriptionStatsCounters> {
+        todo!()
+    }
+
+    async fn get_machines(
+        &self,
+        subscription: &str,
+        start_time: i64,
+        stat_type: Option<SubscriptionMachineState>,
+    ) -> Result<Vec<SubscriptionMachine>> {
+        todo!()
+    }
+}

common/src/database/redis.rs

@@ -1,10 +1,10 @@
 use std::collections::HashMap;
 
-use serde::{ser::SerializeStruct, Serialize, Serializer};
+use serde::{ser::SerializeStruct, Deserialize, Serialize, Serializer};
 
 use crate::{subscription::SubscriptionData, utils, utils::Timestamp};
 
-#[derive(Debug, Serialize, PartialEq, Eq, Clone)]
+#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
 pub struct HeartbeatData {
     machine: String,
     ip: String,