diff --git a/tools/defguard_generator/README.md b/tools/defguard_generator/README.md index 562f45d238..252b7053d8 100644 --- a/tools/defguard_generator/README.md +++ b/tools/defguard_generator/README.md @@ -18,19 +18,19 @@ This means that if you have a development environment set up already it should j ```bash cargo run -p defguard_generator -- vpn-session-stats \ - --location-id 1 \ --num-users 10 \ --devices-per-user 2 \ --sessions-per-device 5 + +# optional: limit generation to one VPN location --location-id ``` ### Session generation logic -For each device the generator always starts with creating an active (not disconnected) session. +For each device the generator always starts with creating an active (not disconnected) session. If there are more sessions per device to be generated it goes backwards in time and creates additional disconnected sessions. Session duration and gaps between sessions are randomized but there is no logic to verify if sessions are overlapping so by default the generator runs a `TRUNCATE` query at the start. To disable this behavior (for example when running it multiple times for separate locations) use the `--no-truncate` CLI flag. - diff --git a/tools/defguard_generator/data/device_names.txt b/tools/defguard_generator/data/device_names.txt new file mode 100644 index 0000000000..42477b4f2d --- /dev/null +++ b/tools/defguard_generator/data/device_names.txt @@ -0,0 +1,39 @@ +Acer Swift +Apple iMac +Apple Mac mini +Apple Mac Pro +Apple Mac Studio +Apple MacBook Air +Apple MacBook Pro +Asus ROG Zephyrus +Asus ZenBook +Chromebook +Dell Alienware +Dell Inspiron +Dell Latitude +Dell OptiPlex +Dell Precision +Dell XPS +Google Pixel +Google Pixel Pro +HP EliteBook +HP Envy +HP Pavilion +HP ProDesk +HP Spectre +iPhone +iPhone Pro +iPhone Pro Max +Lenovo IdeaPad +Lenovo Legion +Lenovo ThinkCentre +Lenovo ThinkPad +Microsoft Surface Laptop +Microsoft Surface Studio +Motorola Edge +OnePlus +Samsung Galaxy A +Samsung Galaxy S +Samsung Galaxy Z Flip +Samsung Galaxy Z Fold +Sony Xperia diff --git a/tools/defguard_generator/data/first_names.txt b/tools/defguard_generator/data/first_names.txt new file mode 100644 index 0000000000..812402c371 --- /dev/null +++ b/tools/defguard_generator/data/first_names.txt @@ -0,0 +1,99 @@ +Elsie +Edward +William +Robert +Annie +Katie +Alma +Thelma +Ellen +Rebecca +Estella +Corine +Margie +Emmie +Myrtis +Pinkie +Myrtie +Jeanette +Rubie +Maud +Trudie +Barbara +Glennie +Lassie +Liza +Hettie +Iona +Lorraine +Willie +Fannie +Sarah +Myrtle +Rosie +Lillian +Nellie +Mamie +Edna +Evelyn +Laura +Mabel +Josie +Pearlie +Rachel +Era +Bettie +Jean +Anne +Lucinda +Dessie +Linnie +Marion +Dovie +Ira +Ossie +Abbie +Clyde +Goldie +Lenora +Aline +Gennie +Vinnie +Exie +Tressie +Joe +Louis +Alfred +Julius +Sidney +Nelson +Isaac +Jimmy +King +Porter +Tommy +Cleo +Dennis +Gerald +Moses +Green +Luke +Richard +Clarence +Ernest +Ralph +Emmett +Dave +Harvey +Ray +Simon +Leslie +Wilson +Emmitt +Gordon +Hollis +Claud +Clayton +Dewey +Lamar +Ruth diff --git a/tools/defguard_generator/data/last_names.txt b/tools/defguard_generator/data/last_names.txt new file mode 100644 index 0000000000..82753cb26a --- /dev/null +++ b/tools/defguard_generator/data/last_names.txt @@ -0,0 +1,100 @@ +Smith +Johnson +Williams +Brown +Jones +Miller +Davis +Wilson +Anderson +Thomas +Taylor +Moore +Jackson +Martin +Thompson +White +Harris +Clark +Lewis +Robinson +Walker +Young +Allen +King +Wright +Scott +Hill +Green +Adams +Nelson +Baker +Hall +Campbell +Mitchell +Carter +Roberts +Phillips +Evans +Turner +Parker +Edwards +Collins +Stewart +Morris +Murphy +Cook +Rogers +Morgan +Cooper +Peterson +Bailey +Reed +Kelly +Howard +Cox +Ward +Richardson +Watson +Brooks +Wood +James +Bennett +Gray +Hughes +Price +Sanders +Myers +Long +Ross +Foster +Powell +Jenkins +Perry +Russell +Sullivan +Bell +Coleman +Henderson +Barnes +Simmons +Patterson +Jordan +Reynolds +Hamilton +Graham +Wallace +Cole +West +Hayes +Woods +Harrison +Gibson +McDonald +Marshall +Murray +Freeman +Wells +Webb +Simpson +Tucker diff --git a/tools/defguard_generator/src/main.rs b/tools/defguard_generator/src/main.rs index a3db8cb6aa..59708fc4ee 100644 --- a/tools/defguard_generator/src/main.rs +++ b/tools/defguard_generator/src/main.rs @@ -34,7 +34,7 @@ enum Commands { /// Generates fake VPN session statistics. VpnSessionStats { #[arg(long)] - location_id: Id, + location_id: Option, #[arg(long)] num_users: usize, #[arg(long)] diff --git a/tools/defguard_generator/src/user_devices.rs b/tools/defguard_generator/src/user_devices.rs index 19b726a2ad..8e71b90873 100644 --- a/tools/defguard_generator/src/user_devices.rs +++ b/tools/defguard_generator/src/user_devices.rs @@ -1,12 +1,16 @@ +use std::collections::HashSet; + use anyhow::Result; use defguard_common::db::{ Id, - models::{Device, User}, + models::{Device, DeviceType, User}, }; use rand::{Rng, rngs::ThreadRng}; use sqlx::PgPool; use tracing::info; +const DEVICE_NAMES: &str = include_str!("../data/device_names.txt"); + pub async fn prepare_user_devices( pool: &PgPool, rng: &mut ThreadRng, @@ -24,14 +28,34 @@ pub async fn prepare_user_devices( ); return Ok(user_devices[..devices_per_user].to_vec()); } + let device_names: Vec<&str> = DEVICE_NAMES.lines().collect(); + let mut taken_names: HashSet = user_devices + .iter() + .map(|device| device.name.clone()) + .collect(); // 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(); + let base_name = device_names[rng.gen_range(0..device_names.len())]; + device.name = unique_device_name(base_name, &mut taken_names); device.user_id = user.id; + device.device_type = DeviceType::User; + device.description = None; let device = device.save(pool).await?; user_devices.push(device); } Ok(user_devices) } + +fn unique_device_name(base: &str, taken: &mut HashSet) -> String { + let mut name = base.to_string(); + let mut suffix = 1; + while taken.contains(&name) { + suffix += 1; + name = format!("{base}{suffix}"); + } + taken.insert(name.clone()); + name +} diff --git a/tools/defguard_generator/src/users.rs b/tools/defguard_generator/src/users.rs index b223b6d120..821da71a4e 100644 --- a/tools/defguard_generator/src/users.rs +++ b/tools/defguard_generator/src/users.rs @@ -1,9 +1,14 @@ +use std::collections::HashSet; + use anyhow::Result; use defguard_common::db::{Id, models::User}; use rand::{Rng, rngs::ThreadRng}; use sqlx::PgPool; use tracing::info; +const FIRST_NAMES: &str = include_str!("../data/first_names.txt"); +const LAST_NAMES: &str = include_str!("../data/last_names.txt"); + pub async fn prepare_users( pool: &PgPool, rng: &mut ThreadRng, @@ -23,12 +28,36 @@ pub async fn prepare_users( return Ok(all_users[..num_users].to_vec()); } + let first_names: Vec<&str> = FIRST_NAMES.lines().collect(); + let last_names: Vec<&str> = LAST_NAMES.lines().collect(); + let mut taken_usernames: HashSet = + all_users.iter().map(|user| user.username.clone()).collect(); + // if there are not enough users create new ones for _ in 0..(num_users - all_users.len()) { - let user: User = rng.r#gen(); + let mut user: User = rng.r#gen(); + let first_name = first_names[rng.gen_range(0..first_names.len())]; + let last_name = last_names[rng.gen_range(0..last_names.len())]; + user.username = unique_username(first_name, last_name, &mut taken_usernames); + user.first_name = first_name.to_string(); + user.last_name = last_name.to_string(); + user.email = format!("{}@defguard.net", user.username); + user.phone = Some("123454321".to_string()); let user = user.save(pool).await?; all_users.push(user); } Ok(all_users) } + +fn unique_username(first_name: &str, last_name: &str, taken: &mut HashSet) -> String { + let base = format!("{first_name}.{last_name}").to_lowercase(); + let mut username = base.clone(); + let mut suffix = 1; + while taken.contains(&username) { + suffix += 1; + username = format!("{base}{suffix}"); + } + taken.insert(username.clone()); + username +} diff --git a/tools/defguard_generator/src/vpn_session_stats.rs b/tools/defguard_generator/src/vpn_session_stats.rs index 21749ed663..6e008a7483 100644 --- a/tools/defguard_generator/src/vpn_session_stats.rs +++ b/tools/defguard_generator/src/vpn_session_stats.rs @@ -12,7 +12,7 @@ use defguard_common::db::{ vpn_session_stats::VpnSessionStats, }, }; -use rand::{Rng, rngs::ThreadRng}; +use rand::{Rng, rngs::ThreadRng, seq::SliceRandom}; use sqlx::{PgConnection, PgPool, QueryBuilder, query}; use tracing::{debug, info}; @@ -23,7 +23,7 @@ const HANDSHAKE_INTERVAL: Duration = Duration::minutes(2); #[derive(Debug)] pub struct VpnSessionGeneratorConfig { - pub location_id: Id, + pub location_id: Option, pub num_users: usize, pub devices_per_user: u8, pub sessions_per_device: u8, @@ -36,7 +36,6 @@ pub async fn generate_vpn_session_stats( config: VpnSessionGeneratorConfig, ) -> Result<()> { info!("Running VPN stats generator with config: {config:#?}"); - let mut rng = rand::thread_rng(); // clear sessions & stats tables unless disabled if !config.no_truncate { @@ -44,17 +43,47 @@ pub async fn generate_vpn_session_stats( truncate_with_restart(&pool).await?; } - // fetch specified location - let location = WireguardNetwork::find_by_id(&pool, config.location_id) - .await? - .expect("Location not found"); + let locations = match config.location_id { + Some(location_id) => { + let location = WireguardNetwork::find_by_id(&pool, location_id) + .await? + .expect("Location not found"); + vec![location] + } + None => WireguardNetwork::all(&pool).await?, + }; + + info!("Generating stats for {} VPN location(s)", locations.len()); + + for location in locations { + generate_stats_for_location(&pool, &config, location).await?; + } + + Ok(()) +} + +async fn generate_stats_for_location( + pool: &PgPool, + config: &VpnSessionGeneratorConfig, + location: WireguardNetwork, +) -> Result<()> { + let mut rng = rand::thread_rng(); + + let location_seed = rng.gen_range(0.0..1_000.0); + + info!( + "Generating VPN stats for location {} ({})", + location.name, location.id + ); // prepare a gateway - let gateway = prepare_gateway(&pool, location.id).await?; + let gateway = prepare_gateway(pool, location.id).await?; // prepare requested number of users - let user_count = config.num_users; - let users = prepare_users(&pool, &mut rng, user_count).await?; + let mut users = prepare_users(pool, &mut rng, config.num_users).await?; + users.shuffle(&mut rng); + let user_count = rng.gen_range((config.num_users / 2).max(1)..=config.num_users.max(1)); + users.truncate(user_count); // generate sessions for each user for (i, user) in users.into_iter().enumerate() { @@ -67,8 +96,8 @@ pub async fn generate_vpn_session_stats( let mut transaction = pool.begin().await?; // prepare requested number of devices - let devices = - prepare_user_devices(&pool, &mut rng, &user, config.devices_per_user as usize).await?; + let device_count = rng.gen_range(1..=config.devices_per_user.max(1)) as usize; + let devices = prepare_user_devices(pool, &mut rng, &user, device_count).await?; let mut used_ips = location.all_used_ips_for_network(&mut transaction).await?; // assign devices to the network if not already assigned @@ -101,8 +130,9 @@ pub async fn generate_vpn_session_stats( // start with the active session let mut session_end = Utc::now().naive_utc(); + let session_count = rng.gen_range(1..=config.sessions_per_device.max(1)); - for i in 0..config.sessions_per_device { + for i in 0..session_count { let session_duration = Duration::minutes(rng.gen_range(10..120)); let session_start = session_end - session_duration; @@ -112,17 +142,16 @@ pub async fn generate_vpn_session_stats( device.id, Some(session_start), None, - ) - .save(&mut *transaction) - .await?; + ); // mark all but the first session as disconnected if i > 0 { session.state = VpnClientSessionState::Disconnected; session.disconnected_at = Some(session_end); - session.save(&mut *transaction).await?; } + let session = session.save(&mut *transaction).await?; + debug!("Created session {session:?}"); generate_mock_session_stats( @@ -133,6 +162,7 @@ pub async fn generate_vpn_session_stats( session_start, session_end, config.stats_batch_size, + location_seed, ) .await?; @@ -171,7 +201,7 @@ async fn prepare_gateway(pool: &PgPool, location_id: Id) -> Result> } } } - +#[allow(clippy::too_many_arguments)] async fn generate_mock_session_stats( transaction: &mut PgConnection, rng: &mut ThreadRng, @@ -180,6 +210,7 @@ async fn generate_mock_session_stats( session_start: NaiveDateTime, session_end: NaiveDateTime, batch_size: u16, + location_seed: f64, ) -> Result<()> { let mut latest_handshake = session_start; let mut next_handshake = latest_handshake + HANDSHAKE_INTERVAL; @@ -190,14 +221,23 @@ async fn generate_mock_session_stats( // assume the IP remains static within a single session let endpoint = random_socket_addr(rng).to_string(); + let upload_scale = rng.gen_range(20_000.0..400_000.0); + let download_scale = rng.gen_range(20_000.0..400_000.0); + // Vector to accumulate stats before batch insertion let mut stats_batch: Vec = Vec::new(); while collected_at <= session_end { + let minutes = collected_at.and_utc().timestamp() as f64 / 60.0; + let activity = |phase: f64| { + (0.5 + 0.35 * (minutes * 0.5 + phase + location_seed).sin()).clamp(0.05, 1.0) + }; + // generate traffic - let upload_diff = rng.gen_range(100..100_000); + let upload_diff = (upload_scale * activity(0.0) * rng.gen_range(0.2..1.8)).max(1.0) as i64; total_upload += upload_diff; - let download_diff = rng.gen_range(100..100_000); + let download_diff = + (download_scale * activity(2.0) * rng.gen_range(0.2..1.8)).max(1.0) as i64; total_download += download_diff; let stats = VpnSessionStats::new( @@ -208,7 +248,7 @@ async fn generate_mock_session_stats( endpoint.clone(), total_upload, total_download, - download_diff, + upload_diff, download_diff, );