use log::trace;
use std::{
sync::{mpsc, Arc},
thread::{self, JoinHandle},
};
use keyhunter::{
ApiKeyCollector, ApiKeyMessage, ApiKeyReceiver, Config, ScriptMessage, WebsiteWalkBuilder,
};
use miette::{Context as _, Error, IntoDiagnostic as _, Result};
#[derive(Debug)]
pub struct Runner {
config: Arc<Config>,
max_walks: usize,
headers: Vec<(String, String)>,
random_ua: bool,
}
impl Runner {
pub fn new(
config: Arc<Config>,
max_walks: usize,
headers: Vec<(String, String)>,
random_ua: bool,
) -> Self {
Self {
config,
max_walks,
headers,
random_ua,
}
}
pub fn run<U: IntoIterator<Item = String> + Send + 'static>(
&self,
urls: U,
) -> (ApiKeyReceiver, JoinHandle<Vec<Error>>) {
let (key_sender, key_receiver) = mpsc::channel::<ApiKeyMessage>();
let config = Arc::clone(&self.config);
let random_ua = self.random_ua;
let walk_builder = WebsiteWalkBuilder::default()
.with_max_walks(self.max_walks)
.with_random_ua(self.random_ua)
.with_headers(self.headers.clone())
.with_shared_cache(true)
.with_cookie_jar(true);
trace!("Starting runner thread");
let handle = thread::spawn(move || {
let mut errors: Vec<Error> = Default::default();
urls.into_iter().for_each(|url| {
let url = if !url.contains("://") {
String::from("https://") + &url
} else {
url
};
let (tx_scripts, rx_scripts) = mpsc::channel::<ScriptMessage>();
let walker = walk_builder.build(tx_scripts.clone());
let collector =
ApiKeyCollector::new(config.clone(), rx_scripts, key_sender.clone())
.with_random_ua(random_ua);
let moved_url = url.clone();
let walk_handle = thread::spawn(move || {
let result = walker.walk(&moved_url);
if let Err(ref err) = result {
println!("{:?}", err);
tx_scripts
.send(ScriptMessage::Done)
.into_diagnostic()
.context("Failed to send stop signal over script channel")
.unwrap();
}
result
});
trace!("Starting API key collector thread");
let collector_handle = thread::spawn(move || collector.collect());
if let Err(error) = Self::join(&url, collector_handle, walk_handle) {
errors.push(error)
}
});
debug!("Sending stop signal to API key collector");
key_sender
.send(ApiKeyMessage::Stop)
.into_diagnostic()
.context("Failed to close API key channel")
.unwrap();
debug!("Scraping completed");
errors
});
(key_receiver, handle)
}
fn join<S: AsRef<str>>(
url: S,
collector_handle: JoinHandle<()>,
walk_handle: JoinHandle<Result<()>>,
) -> Result<()> {
let _url = url.as_ref();
debug!("Waiting for collector thread to finish...");
collector_handle
.join()
.expect("ApiKeyCollector thread should have joined successfully");
debug!("Waiting for walker thread to finish...");
walk_handle
.join()
.expect("WebsiteWalker thread should have joined successfully")
}
}