Skip to content

Commit 329c412

Browse files
CyberdolfiCyberdolfi
committed
feat(rescan)!: add rescan system
1 parent a58b5d7 commit 329c412

5 files changed

Lines changed: 177 additions & 24 deletions

File tree

‎src/database/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ pub mod server;
1010

1111
// Database Server
1212

13-
#[derive(Debug, Clone, Serialize, Deserialize)]
13+
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
1414
pub struct ServerInfo {
1515
pub server_id: Option<i64>, // Key
1616
pub server_ip: String,

‎src/database/server.rs‎

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use sqlx::{Postgres, Transaction};
1+
use sqlx::{Postgres, Row, Transaction};
22
use crate::database::{parse_players, pool, Player, PlayerHistory, ServerHistory, ServerInfo};
33

44
pub async fn insert_servers(results: &Vec<(ServerInfo, ServerHistory)>) -> Result<(), sqlx::Error> {
@@ -89,23 +89,23 @@ pub async fn insert_players(player_data: &Vec<(Player, PlayerHistory)>, tx: &mut
8989
Ok(())
9090
}
9191

92-
pub async fn get_total_servers() -> Result<i64, sqlx::Error> {
93-
let (count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM servers")
94-
.fetch_one(pool::get_pool())
95-
.await?;
96-
Ok(count)
97-
}
92+
pub async fn get_total_servers() -> Result<Vec<ServerInfo>, sqlx::Error> {
93+
let pool = pool::get_pool();
9894

99-
pub async fn get_total_history() -> Result<i64, sqlx::Error> {
100-
let (count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM server_history")
101-
.fetch_one(pool::get_pool())
95+
let rows = sqlx::query("SELECT server_id, server_ip, server_port, last_seen, discovered, bedrock, country FROM servers")
96+
.fetch_all(pool)
10297
.await?;
103-
Ok(count)
104-
}
10598

106-
pub async fn get_total_player_history() -> Result<i64, sqlx::Error> {
107-
let (count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM player_history")
108-
.fetch_one(pool::get_pool())
109-
.await?;
110-
Ok(count)
99+
let servers = rows.into_iter().map(|row| {
100+
ServerInfo {
101+
server_id: Some(row.get::<i32, _>("server_id") as i64),
102+
server_ip: row.get::<String, _>("server_ip"),
103+
server_port: row.get::<i32, _>("server_port") as u16,
104+
last_seen: row.get::<i64, _>("last_seen"),
105+
discovered: row.get::<i64, _>("discovered"),
106+
bedrock: row.get::<bool, _>("bedrock"),
107+
country: row.get::<Option<String>, _>("country"),
108+
}
109+
}).collect();
110+
Ok(servers)
111111
}

‎src/main.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,9 @@ use tokio::io::{AsyncWriteExt, BufWriter};
1111
use crate::config::{DatabaseConfig, MainConfig};
1212
use crate::logger::DefaultColor;
1313
use crate::manager::TaskManager;
14-
use crate::randomizer::{IpGenerator, IpGeneratorBuilder, IpType};
14+
use crate::randomizer::{IpGenerator, IpType};
1515
use crate::scanning::crawler;
16+
use crate::scanning::rescanner::rescan;
1617
use crate::updater::GithubAPI;
1718

1819
mod updater;
@@ -25,7 +26,6 @@ mod database;
2526
mod config;
2627
mod cli;
2728
mod scanning;
28-
mod webapi;
2929

3030
type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
3131

@@ -244,9 +244,9 @@ async fn main() -> Result<()> {
244244
}
245245
}).await;
246246
},
247-
248-
cli::Commands::Start { ip, port } => {
249-
webapi::start(port).await;
247+
248+
cli::Commands::Rescan => {
249+
rescan().await;
250250
}
251251
}
252252

‎src/scanning/mod.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
pub mod scanner;
22
pub mod file_scanner;
33
pub mod utils;
4-
pub mod crawler;
4+
pub mod crawler;
5+
pub mod rescanner;

‎src/scanning/rescanner.rs‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
use std::net::Ipv4Addr;
2+
use std::str::FromStr;
3+
use std::time::{Duration, Instant};
4+
use colored_text::Colorize;
5+
use futures::StreamExt;
6+
use crate::database::{parse_server, server, ServerHistory, ServerInfo};
7+
use crate::logger;
8+
use crate::config::MainConfig;
9+
use crate::logger::DefaultColor;
10+
use crate::manager::TaskManager;
11+
use crate::scanning::scanner::{scan, ScanConfig};
12+
use crate::scanning::utils::{format_time, prettier_ping_result, save_server};
13+
14+
pub async fn rescan() {
15+
let _ = TaskManager::spawn("Rescan", move |cancel_token| async move {
16+
logger::info("Getting all servers from the database...".into()).send().await;
17+
18+
let servers = match server::get_total_servers().await {
19+
Ok(s) => s,
20+
Err(e) => {
21+
logger::critical(format!("Failed to load servers: {}", e.hex(DefaultColor::Highlight.hex()))).prefix("Rescan").send().await;
22+
return;
23+
}
24+
};
25+
26+
if servers.is_empty() {
27+
logger::warning("No servers found in database to rescan.".into()).prefix("Rescan").send().await;
28+
return;
29+
}
30+
31+
32+
let start_time = Instant::now();
33+
34+
let mut targets = Vec::new();
35+
for s in servers {
36+
if let Ok(ipv4) = Ipv4Addr::from_str(&s.server_ip) {
37+
targets.push((ipv4, s.server_port));
38+
}
39+
}
40+
41+
let mut found_batch: Vec<(ServerInfo, ServerHistory)> = Vec::new();
42+
43+
let total_targets = targets.len();
44+
let mut total_found_count = 0;
45+
let mut processed_count = 0;
46+
47+
logger::info(format!(
48+
"Scanning {} targets...",
49+
total_targets.hex(DefaultColor::Highlight.hex())
50+
)).prefix("Rescan").send().await;
51+
52+
let main_cfg = MainConfig::get().expect("Config not loaded!");
53+
54+
let config = ScanConfig {
55+
ping_timeout: Duration::from_millis(main_cfg.general.ping_timeout),
56+
query_timeout: Duration::from_millis(main_cfg.general.query_timeout),
57+
join_timeout: Duration::from_millis(main_cfg.general.join_timeout),
58+
with_uuid: main_cfg.general.do_uuid_fetch,
59+
max_tasks: main_cfg.get_scanner_tasks(),
60+
..ScanConfig::default()
61+
};
62+
63+
// Core part: scanning
64+
let scan_stream = scan(targets, config);
65+
tokio::pin!(scan_stream);
66+
67+
// Scan stream
68+
while let Some(maybe_result) = scan_stream.next().await {
69+
if cancel_token.is_cancelled() {
70+
logger::warning("Scan interrupted. Saving results...".to_string())
71+
.prefix("Rescan").send().await;
72+
break;
73+
}
74+
75+
processed_count += 1;
76+
77+
// Success
78+
if let Some(result) = maybe_result {
79+
let parsed = parse_server(result.ip, result.port, result.ping.clone(), result.query, result.join);
80+
found_batch.push(parsed);
81+
total_found_count += 1;
82+
83+
let mut output = String::new();
84+
output.push_str(
85+
&format!(
86+
"Found server: {}:{}\n",
87+
result.ip.to_string().hex(DefaultColor::Highlight.hex()),
88+
result.port.hex(DefaultColor::Highlight.hex())
89+
)
90+
);
91+
output.push_str(&prettier_ping_result(result.ping).await);
92+
logger::success(output).prefix("Rescan").send().await;
93+
94+
if found_batch.len() >= 30 {
95+
let batch_to_insert = std::mem::take(&mut found_batch);
96+
save_server(&batch_to_insert).await;
97+
}
98+
}
99+
100+
// Progress calc
101+
let elapsed = start_time.elapsed().as_secs_f64();
102+
let ips_per_second = processed_count as f64 / elapsed;
103+
104+
if ips_per_second > 0.0 {
105+
let remaining_ips = total_targets.saturating_sub(processed_count);
106+
let remaining_secs = remaining_ips as f64 / ips_per_second;
107+
let percent = format!("{:.2}", (processed_count as f64 / total_targets as f64) * 100.0);
108+
109+
if processed_count % 10000 == 0 || processed_count == total_targets {
110+
logger::info(format!(
111+
"Progress: {}/{} IPs ({}%) - ETA: {}",
112+
processed_count.hex(DefaultColor::Highlight.hex()),
113+
total_targets.hex(DefaultColor::Highlight.hex()),
114+
percent.hex(DefaultColor::Highlight.hex()),
115+
format_time(remaining_secs as u64)
116+
)).prefix("Rescan").send().await;
117+
}
118+
}
119+
}
120+
121+
if !found_batch.is_empty() {
122+
save_server(&found_batch).await;
123+
}
124+
125+
// Finished
126+
let elapsed_time = start_time.elapsed();
127+
128+
let pps = if elapsed_time.as_secs() > 0 {
129+
total_targets as f64 / elapsed_time.as_secs_f64()
130+
} else {
131+
0.0
132+
};
133+
134+
let percent = if total_found_count > 0 {
135+
format!("{:.2}", (total_found_count as f64 / processed_count as f64) * 100.0)
136+
} else {
137+
"0.00".to_string()
138+
};
139+
140+
logger::info(
141+
format!(
142+
"Rescan finished in {}. Updated {} servers from {} targets, {}%. ({}{})",
143+
format_time(elapsed_time.as_secs()).hex(DefaultColor::Highlight.hex()),
144+
total_found_count.hex(DefaultColor::Highlight.hex()),
145+
total_targets.hex(DefaultColor::Highlight.hex()),
146+
percent.hex(DefaultColor::Highlight.hex()),
147+
pps.round().hex(DefaultColor::Highlight.hex()),
148+
"pps".hex(DefaultColor::DarkHighlight.hex())
149+
)
150+
).send().await;
151+
}).await;
152+
}

0 commit comments

Comments
 (0)