firststep_name_lib/
lib.rs

1use colored::*;
2use futures_util::sink::SinkExt;
3use poem::web::websocket::{Message, WebSocketStream};
4use reqwest::Client;
5use serde::{Deserialize, Serialize};
6use std::error::Error;
7use std::fs::File;
8use std::fs::OpenOptions;
9use std::io::Write;
10use std::sync::Arc;
11use std::time::Duration;
12use tokio::sync::Mutex;
13use tokio::time::sleep;
14use url::Url;
15
16const DEFAULT_DATA_URL: &str =
17    "https://raw.githubusercontent.com/WebBreacher/WhatsMyName/main/wmn-data.json";
18
19#[derive(Serialize, Deserialize, Debug, Clone)]
20pub struct ProgressUpdate {
21    site: String,
22    status: String,
23    url: String,
24    logo_url: String,
25    error: Option<String>,
26    is_taken: bool,
27    completed: usize,
28    total: usize,
29}
30
31#[derive(Serialize, Deserialize, Debug, Clone)]
32pub struct SiteData {
33    name: String,
34    uri_check: String,
35    e_code: u16,
36    e_string: String,
37    m_string: String,
38    m_code: u16,
39    known: Vec<String>,
40    cat: String,
41}
42
43#[derive(Serialize, Deserialize, Debug)]
44pub struct SitesFile {
45    license: Vec<String>,
46    authors: Vec<String>,
47    categories: Vec<String>,
48    pub sites: Vec<SiteData>,
49}
50
51#[derive(Serialize, Deserialize, Debug, Clone)]
52pub struct CheckResult {
53    site: String,
54    status: String,
55    url: String,
56    logo_url: String,
57    error: Option<String>,
58}
59
60#[derive(Serialize, Deserialize, Debug)]
61struct Report {
62    username: String,
63    generated_at: String,
64    results: Vec<CheckResult>,
65}
66
67pub fn get_site_logo(domain_name: &str) -> String {
68    match domain_name {
69        "t.me" => "https://logo.clearbit.com/telegram.org".to_string(),
70        "giters.com" => "https://giters.com/images/favicon.svg".to_string(),
71        "ko-fi.com" => {
72            "https://storage.ko-fi.com/cdn/brandasset/kofi_s_logo_nolabel.png".to_string()
73        }
74        _ => format!("https://logo.clearbit.com/{}", domain_name),
75    }
76}
77
78pub fn extract_domain(url_str: &str) -> Option<String> {
79    if let Ok(url) = Url::parse(url_str) {
80        if let Some(host) = url.host_str() {
81            // Get the base domain (example.com from subdomain.example.com)
82            let parts: Vec<&str> = host.split('.').collect();
83            if parts.len() >= 2 {
84                // For most domains, return the last two parts
85                return Some(format!(
86                    "{}.{}",
87                    parts[parts.len() - 2],
88                    parts[parts.len() - 1]
89                ));
90            } else {
91                return Some(host.to_string());
92            }
93        }
94    }
95    None
96}
97
98pub async fn download_sites_data(client: &Client, output_file: &str) -> Result<(), Box<dyn Error>> {
99    println!("Downloading sites data from {}...", DEFAULT_DATA_URL);
100
101    let response = client.get(DEFAULT_DATA_URL).send().await?;
102
103    if response.status().is_success() {
104        let data = response.text().await?;
105        let mut file = OpenOptions::new()
106            .write(true)
107            .create(true)
108            .truncate(true)
109            .open(output_file)?;
110
111        file.write_all(data.as_bytes())?;
112        println!("Successfully downloaded sites data to {}", output_file);
113        Ok(())
114    } else {
115        Err(format!("Failed to download data: HTTP {}", response.status()).into())
116    }
117}
118
119use futures_util::stream::SplitSink;
120
121pub async fn check_username_from_webserver(
122    client: &Client,
123    username: &str,
124    sites_data: &[SiteData],
125    _threads: usize, // Ignored parameter, but kept for compatibility
126    ws_sink: Option<Arc<Mutex<SplitSink<WebSocketStream, Message>>>>,
127) -> Vec<CheckResult> {
128    println!("Checking availability for username: {}\n", username);
129
130    let total = sites_data.len();
131    let mut all_results = Vec::new();
132    let mut completed = 0;
133
134    for site in sites_data {
135        let uri_string = site.uri_check.replace("{account}", username);
136        let domain = extract_domain(&uri_string).unwrap_or_else(|| "unknown.com".to_string());
137        let logo_url = get_site_logo(&domain);
138
139        let (is_taken, status, error) = match check_site(client, site, &uri_string).await {
140            Ok((is_taken, status)) => (is_taken, status, None),
141            Err(e) => (false, "Error".to_string(), Some(e.to_string())),
142        };
143
144        completed += 1;
145
146        if let Some(ws_sink) = &ws_sink {
147            let update = ProgressUpdate {
148                site: site.name.clone(),
149                status: status.clone(),
150                url: uri_string.clone(),
151                logo_url: logo_url.clone(),
152                error: error.clone(),
153                is_taken,
154                completed,
155                total,
156            };
157            if let Ok(json) = serde_json::to_string(&update) {
158                let mut sink = ws_sink.lock().await;
159                if sink.send(Message::Text(json.into())).await.is_err() {
160                    eprintln!("Failed to send WebSocket message");
161                }
162            }
163        }
164
165        let color = if is_taken { "red" } else { "green" };
166        println!("{} {} - {}", status.color(color), site.name, uri_string);
167
168        all_results.push(CheckResult {
169            site: site.name.clone(),
170            status,
171            url: uri_string,
172            logo_url,
173            error,
174        });
175
176        sleep(Duration::from_millis(100)).await;
177    }
178
179    if let Some(ws_sink) = &ws_sink {
180        let completion_msg = serde_json::json!({
181            "completed": true,
182            "total": total
183        })
184        .to_string();
185        let mut sink = ws_sink.lock().await;
186        if sink
187            .send(Message::Text(completion_msg.into()))
188            .await
189            .is_err()
190        {
191            eprintln!("Failed to send completion message");
192        }
193    }
194
195    all_results
196}
197
198pub async fn check_username(
199    client: &Client,
200    username: &str,
201    sites_data: &[SiteData],
202    threads: usize,
203) -> Vec<CheckResult> {
204    println!("Checking availability for username: {}\n", username);
205
206    // Create a shared client for concurrent requests
207    let client = client.clone();
208
209    // The maximum number of concurrent requests
210    let concurrent_limit = threads;
211
212    // Process sites in chunks to control concurrency
213    let mut all_results = Vec::new();
214
215    // Process sites in chunks to avoid overwhelming APIs
216    for chunk in sites_data.chunks(concurrent_limit) {
217        let mut tasks = Vec::new();
218
219        // Start tasks for each site in the chunk
220        for site in chunk {
221            let site = site.clone();
222            let client = client.clone();
223            let username = username.to_string();
224
225            // Spawn an async task for each site check
226            let task = tokio::spawn(async move {
227                let uri_string = site.uri_check.replace("{account}", &username);
228                let domain =
229                    extract_domain(&uri_string).unwrap_or_else(|| "unknown.com".to_string());
230                let logo_url = get_site_logo(&domain);
231
232                match check_site(&client, &site, &uri_string).await {
233                    Ok((is_taken, status)) => (
234                        site.name.clone(),
235                        status,
236                        uri_string,
237                        logo_url,
238                        None,
239                        is_taken,
240                    ),
241                    Err(e) => (
242                        site.name.clone(),
243                        "Error".to_string(),
244                        uri_string,
245                        logo_url,
246                        Some(e.to_string()),
247                        false,
248                    ),
249                }
250            });
251
252            tasks.push(task);
253        }
254
255        // Wait for all tasks in this chunk to complete
256        for task in tasks {
257            if let Ok((site_name, status, url, logo_url, error, is_taken)) = task.await {
258                let color = if is_taken { "red" } else { "green" };
259                if let Some(err) = &error {
260                    println!("{} {} - {}", "Error".color("yellow"), site_name, err);
261                } else {
262                    println!("{} {} - {}", status.color(color), site_name, url);
263                }
264
265                all_results.push(CheckResult {
266                    site: site_name,
267                    status,
268                    url,
269                    logo_url,
270                    error,
271                });
272            }
273        }
274
275        // Add a small delay between chunks to be nice to APIs
276        sleep(Duration::from_millis(200)).await;
277    }
278
279    all_results
280}
281
282async fn check_site(
283    client: &Client,
284    site: &SiteData,
285    uri: &str,
286) -> Result<(bool, String), Box<dyn Error>> {
287    let response = client
288        .get(uri)
289        .timeout(Duration::from_secs(10))
290        .send()
291        .await?;
292
293    let status = response.status();
294    let body = response.text().await?;
295
296    let is_taken = if status.as_u16() == site.e_code && body.contains(&site.e_string) {
297        true
298    } else if status.as_u16() == site.m_code && body.contains(&site.m_string) {
299        // If we get the "missing" code and string, the user does not exist
300        false
301    } else {
302        // In all other cases, assume the user is available
303        false
304    };
305
306    let status_text = if is_taken {
307        "Taken".to_string()
308    } else {
309        "Available".to_string()
310    };
311
312    // Here you put the code for real-time status update if needed
313
314    Ok((is_taken, status_text))
315}
316
317pub fn save_txt_report(username: &str, results: &[CheckResult]) -> Result<(), Box<dyn Error>> {
318    let filename = format!("{}_report.txt", username);
319    let mut file = OpenOptions::new()
320        .write(true)
321        .create(true)
322        .truncate(true)
323        .open(filename)?;
324
325    writeln!(file, "Username availability report for: {}", username)?;
326    writeln!(file, "Generated on: {}", chrono::Local::now())?;
327    writeln!(file, "{}", "-".repeat(80))?;
328
329    for result in results {
330        writeln!(file, "{}: {}", result.site, result.status)?;
331        writeln!(file, "URL: {}", result.url)?;
332        writeln!(file, "Logo: {}", result.logo_url)?;
333        if let Some(error) = &result.error {
334            writeln!(file, "Error: {}", error)?;
335        }
336        writeln!(file, "{}", "-".repeat(40))?;
337    }
338
339    Ok(())
340}
341
342pub fn save_json_report(username: &str, results: &[CheckResult]) -> Result<(), Box<dyn Error>> {
343    let filename = format!("{}_report.json", username);
344    let report = Report {
345        username: username.to_string(),
346        generated_at: chrono::Local::now().to_string(),
347        results: results.to_vec(),
348    };
349
350    let file = File::create(filename)?;
351    serde_json::to_writer_pretty(file, &report)?;
352
353    Ok(())
354}
firststep_name_lib/lib.rs

firststep_name_lib/
lib.rs
