add VPN sessions & stats generator

Cargo.toml

@@ -6,7 +6,8 @@ repository = "https://github.com/DefGuard/defguard"
 rust-version = "1.85.1"
 
 [workspace]
-members = ["crates/*"]
+members = ["crates/*", "tools/*"]
+default-members = ["crates/*"]
 resolver = "2"
 
 [workspace.dependencies]

crates/defguard_common/src/db/models/device.rs

@@ -106,10 +106,10 @@ impl fmt::Display for Device<Id> {
     }
 }
 
-impl Distribution<Device<Id>> for Standard {
-    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Device<Id> {
+impl Distribution<Device> for Standard {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Device {
         Device {
-            id: rng.r#gen(),
+            id: NoId,
             name: Alphanumeric.sample_string(rng, 8),
             wireguard_pubkey: Alphanumeric.sample_string(rng, 32),
             user_id: rng.r#gen(),

crates/defguard_core/src/enterprise/firewall/tests.rs

@@ -51,8 +51,8 @@ fn random_user_with_id<R: Rng>(rng: &mut R, id: Id) -> User<Id> {
 }
 
 fn random_network_device_with_id<R: Rng>(rng: &mut R, id: Id) -> Device<Id> {
-    let mut device: Device<Id> = rng.r#gen();
-    device.id = id;
+    let device: Device = rng.r#gen();
+    let mut device = device.with_id(id);
     device.device_type = DeviceType::Network;
     device
 }

deny.toml

@@ -161,6 +161,10 @@ exceptions = [
         "AGPL-3.0-only",
         "AGPL-3.0-or-later",
     ], crate = "defguard_certs" },
+    { allow = [
+        "AGPL-3.0-only",
+        "AGPL-3.0-or-later",
+    ], crate = "defguard_generator" },
 ]
 
 # Some crates don't have (easily) machine readable licensing information,

tools/defguard_generator/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "defguard_generator"
+version = "0.0.0"
+edition = "2024"
+license-file = "../../LICENSE.md"
+homepage = "https://defguard.net/"
+repository = "https://github.com/DefGuard/defguard"
+rust-version = "1.85.1"
+
+[dependencies]
+defguard_common = { workspace = true }
+chrono = { workspace = true }
+rand = { workspace = true }
+clap = { workspace = true, features = ["derive"] }
+sqlx = { workspace = true, features = ["postgres", "runtime-tokio-native-tls", "chrono"] }
+tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
+anyhow = { workspace = true }
+tracing = { workspace = true }
+tracing-subscriber = { workspace = true, features = ["env-filter"] }
+
+[dev-dependencies]

tools/defguard_generator/README.md

@@ -0,0 +1,14 @@
+# Defguard Object Generator
+
+This crate contains a simple generator for creating users, devices, stats etc during development.
+
+### Usage
+
+```bash
+cargo run -p defguard_generator -- vpn-session-stats \
+    --location-id 1 \
+    --num-users 10 \
+    --devices-per-user 2 \
+    --sessions-per-device 5
+```
+

tools/defguard_generator/src/lib.rs

@@ -0,0 +1,3 @@
+pub mod user_devices;
+pub mod users;
+pub mod vpn_session_stats;

tools/defguard_generator/src/main.rs

@@ -0,0 +1,84 @@
+use anyhow::Result;
+use clap::{Parser, Subcommand};
+use defguard_common::db::{Id, init_db};
+use defguard_generator::vpn_session_stats::{
+    VpnSessionGeneratorConfig, generate_vpn_session_stats,
+};
+use tracing::Level;
+
+#[derive(Parser)]
+#[command(about, long_about = None)]
+struct Cli {
+    #[arg(long, env = "DEFGUARD_DB_HOST", default_value = "localhost")]
+    pub database_host: String,
+
+    #[arg(long, env = "DEFGUARD_DB_PORT", default_value_t = 5432)]
+    pub database_port: u16,
+
+    #[arg(long, env = "DEFGUARD_DB_NAME", default_value = "defguard")]
+    pub database_name: String,
+
+    #[arg(long, env = "DEFGUARD_DB_USER", default_value = "defguard")]
+    pub database_user: String,
+
+    #[arg(long, env = "DEFGUARD_DB_PASSWORD", default_value = "")]
+    pub database_password: String,
+
+    #[command(subcommand)]
+    command: Commands,
+}
+
+#[derive(Subcommand)]
+enum Commands {
+    /// generates VPN session stats
+    VpnSessionStats {
+        #[arg(long)]
+        location_id: Id,
+        #[arg(long)]
+        num_users: u16,
+        #[arg(long)]
+        devices_per_user: u8,
+        #[arg(long)]
+        sessions_per_device: u8,
+    },
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    // Initialize logging
+    tracing_subscriber::fmt().with_max_level(Level::INFO).init();
+
+    // parse CLI options
+    let cli = Cli::parse();
+
+    // setup DB pool
+    let pool = init_db(
+        &cli.database_host,
+        cli.database_port,
+        &cli.database_name,
+        &cli.database_user,
+        &cli.database_password,
+    )
+    .await;
+
+    // execute based on the selected subcommand
+    match cli.command {
+        Commands::VpnSessionStats {
+            location_id,
+            num_users,
+            devices_per_user,
+            sessions_per_device,
+        } => {
+            let config = VpnSessionGeneratorConfig {
+                location_id,
+                num_users,
+                devices_per_user,
+                sessions_per_device,
+            };
+
+            generate_vpn_session_stats(pool, config).await?;
+        }
+    };
+
+    Ok(())
+}

tools/defguard_generator/src/user_devices.rs

@@ -0,0 +1,37 @@
+use anyhow::Result;
+use defguard_common::db::{
+    Id,
+    models::{Device, User},
+};
+use rand::{Rng, rngs::ThreadRng};
+use sqlx::PgPool;
+use tracing::info;
+
+pub async fn prepare_user_devices(
+    pool: &PgPool,
+    rng: &mut ThreadRng,
+    user: &User<Id>,
+    devices_per_user: usize,
+) -> Result<Vec<Device<Id>>> {
+    // fetch all existing devices for a given user
+    let mut user_devices = Device::all_for_username(pool, &user.username).await?;
+
+    // if there are enough users just return the required number
+    if user_devices.len() >= devices_per_user {
+        info!(
+            "Found {} existing devices for user {user} in the database. Using the required number.",
+            user_devices.len()
+        );
+        return Ok(user_devices[..devices_per_user].to_vec());
+    }
+
+    // if there are not enough users create new ones
+    for _ in 0..(devices_per_user - user_devices.len()) {
+        let mut device: Device = rng.r#gen();
+        device.user_id = user.id;
+        let device = device.save(pool).await?;
+        user_devices.push(device);
+    }
+
+    Ok(user_devices)
+}

tools/defguard_generator/src/users.rs

@@ -0,0 +1,34 @@
+use anyhow::Result;
+use defguard_common::db::{Id, models::User};
+use rand::{Rng, rngs::ThreadRng};
+use sqlx::PgPool;
+use tracing::info;
+
+pub async fn prepare_users(
+    pool: &PgPool,
+    rng: &mut ThreadRng,
+    num_users: usize,
+) -> Result<Vec<User<Id>>> {
+    info!("Preparing {num_users} random users for generating VPN session stats");
+
+    // fetch all existing users
+    let mut all_users = User::all(pool).await?;
+
+    // if there are enough users just return the required number
+    if all_users.len() >= num_users {
+        info!(
+            "Found {} existing users in the database. Using the required number.",
+            all_users.len()
+        );
+        return Ok(all_users[..num_users].to_vec());
+    }
+
+    // if there are not enough users create new ones
+    for _ in 0..(num_users - all_users.len()) {
+        let user: User = rng.r#gen();
+        let user = user.save(pool).await?;
+        all_users.push(user);
+    }
+
+    Ok(all_users)
+}