YGGverse/aquatic-crawler
1
0
Fork 0
mirror of https://github.com/YGGverse/aquatic-crawler.git synced 2026-03-31 17:15:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

implement some 0.2.0 features

Browse source
This commit is contained in:
yggverse 2025-06-14 23:33:11 +03:00
parent 18116491cb
commit 10c2d7e855
11 changed files with 319 additions and 157 deletions
Download patch file Download diff file Expand all files Collapse all files

31
src/argument.rs
View file

@ -7,7 +7,7 @@ pub struct Argument {
///
/// * `e` - error
/// * `i` - info
/// * `t` - trace (e.g. to run with `RUST_LOG=librqbit=trace`)
/// * `t` - trace (run with `RUST_LOG=librqbit=trace`)
#[arg(short, long, default_value_t = String::from("ei"))]
pub debug: String,
@ -21,9 +21,9 @@ pub struct Argument {
#[arg(long)]
pub infohash_file: Vec<String>,
/// Directory path to store reload data (e.g. `.torrent` files)
/// Directory path to store preloaded data (e.g. `.torrent` files)
#[arg(long)]
pub storage: Option<String>,
pub storage: String,
/// Define custom tracker(s) to preload the `.torrent` files info
#[arg(long)]
@ -45,6 +45,23 @@ pub struct Argument {
#[arg(long, default_value_t = false)]
pub enable_upload: bool,
/// Preload files match regex pattern (list only without preload by default)
///
/// ## Example:
///
/// Filter by image ext
/// ```
/// --preload-regex '\.(png|gif|jpeg|webp)$'
/// ```
///
/// * requires `storage` argument defined
#[arg(long)]
pub preload_regex: Option<String>,
/// Save resolved torrent files to the `storage` location
#[arg(long, default_value_t = true)]
pub save_torrents: bool,
/// Use `socks5://[username:password@]host:port`
#[arg(long)]
pub socks_proxy_url: Option<String>,
@ -52,4 +69,12 @@ pub struct Argument {
/// Crawl loop delay in seconds
#[arg(short, default_value_t = 300)]
pub sleep: u64,
/// Limit upload speed (b/s)
#[arg(long)]
pub upload_limit: Option<u32>,
/// Limit download speed (b/s)
#[arg(long)]
pub download_limit: Option<u32>,
}

1
src/database.rs
View file

@ -1 +0,0 @@
pub mod torrent;

38
src/database/torrent.rs
View file

@ -1,38 +0,0 @@
use anyhow::{Result, bail};
use std::{fs, io::Write, path::PathBuf, str::FromStr};
pub struct Storage(PathBuf);
impl Storage {
pub fn init(storage: &str, clear: bool) -> Result<Self> {
let p = PathBuf::from_str(storage)?;
if let Ok(t) = fs::metadata(&p) {
if t.is_file() {
bail!("Target destination is not directory!")
}
if t.is_dir() && clear {
fs::remove_dir_all(&p)?;
}
}
fs::create_dir_all(&p)?;
Ok(Self(p))
}
pub fn exists(&self, infohash: &str) -> bool {
fs::metadata(self.filename(infohash)).is_ok_and(|p| p.is_file())
}
pub fn save(&self, infohash: &str, bytes: &[u8]) -> Result<PathBuf> {
let p = self.filename(infohash);
let mut f = fs::File::create(&p)?;
f.write_all(bytes)?;
Ok(p)
}
fn filename(&self, infohash: &str) -> PathBuf {
let mut p = PathBuf::new();
p.push(&self.0);
p.push(format!("{infohash}.torrent"));
p
}
}

29
src/debug.rs
View file

@ -1,9 +1,28 @@
pub fn error(e: &anyhow::Error) {
eprintln!("[{}] [error] {e}", now())
}
mod level;
use level::Level;
pub fn info(message: String) {
println!("[{}] [info] {message}", now())
pub struct Debug(Vec<Level>);
impl Debug {
pub fn init(levels: &str) -> anyhow::Result<Self> {
let mut l = Vec::with_capacity(levels.len());
for s in levels.to_lowercase().chars() {
l.push(Level::parse(s)?);
}
Ok(Self(l))
}
pub fn error(&self, message: &str) {
if self.0.contains(&Level::Error) {
eprintln!("[{}] [error] {message}", now());
}
}
pub fn info(&self, message: &str) {
if self.0.contains(&Level::Info) {
println!("[{}] [info] {message}", now());
}
}
}
fn now() -> u128 {

22
src/debug/level.rs Normal file
View file

@ -0,0 +1,22 @@
use anyhow::{Result, bail};
#[derive(PartialEq)]
pub enum Level {
Error,
Info,
Trace,
}
impl Level {
pub fn parse(value: char) -> Result<Self> {
match value {
'e' => Ok(Self::Error),
'i' => Ok(Self::Info),
't' => {
tracing_subscriber::fmt::init();
Ok(Self::Trace)
}
_ => bail!("Unsupported debug value `{value}`!"),
}
}
}

213
src/main.rs
View file

@ -1,132 +1,141 @@
mod api;
mod argument;
mod database;
mod debug;
mod peers;
mod storage;
mod trackers;
use anyhow::Result;
use debug::Debug;
use storage::Storage;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
async fn main() -> Result<()> {
use clap::Parser;
use librqbit::SessionOptions;
use std::str::FromStr;
use librqbit::{AddTorrent, AddTorrentOptions, AddTorrentResponse, SessionOptions};
use std::{num::NonZero, time::Duration};
let argument = argument::Argument::parse();
// calculate debug level once
let is_debug_i = argument.debug.contains("i");
let is_debug_e = argument.debug.contains("e");
if argument.debug.contains("t") {
tracing_subscriber::fmt::init()
}
// init shared members
let torrent_storage = if let Some(t) = argument.storage {
let s = database::torrent::Storage::init(&t, argument.clear)?;
if argument.clear && is_debug_i {
debug::info(String::from("Cleanup torrent storage"));
}
Some(s)
} else {
None
};
let mut trackers = std::collections::HashSet::with_capacity(argument.torrent_tracker.len());
for tracker in argument.torrent_tracker {
trackers.insert(url::Url::from_str(&tracker)?);
}
let mut peers = Vec::with_capacity(argument.initial_peer.len());
for peer in argument.initial_peer {
peers.push(std::net::SocketAddr::from_str(&peer)?);
}
// init components
let arg = argument::Argument::parse();
let debug = Debug::init(&arg.debug)?;
let peers = peers::Peers::init(&arg.initial_peer)?;
let storage = Storage::init(&arg.storage, arg.clear)?;
let trackers = trackers::Trackers::init(&arg.torrent_tracker)?;
let session = librqbit::Session::new_with_opts(
storage.path(),
SessionOptions {
disable_upload: !arg.enable_upload,
disable_dht: !arg.enable_dht,
disable_dht_persistence: true,
persistence: None,
ratelimits: librqbit::limits::LimitsConfig {
upload_bps: arg.upload_limit.and_then(NonZero::new),
download_bps: arg.download_limit.and_then(NonZero::new),
},
trackers: trackers.clone(),
..SessionOptions::default()
},
)
.await?;
// begin
if is_debug_i {
debug::info(String::from("Crawler started"));
}
debug.info("Crawler started");
loop {
if is_debug_i {
debug::info(String::from("New index session begin..."));
}
debug.info("Index queue begin...");
let mut total = 0;
let session = librqbit::Session::new_with_opts(
std::path::PathBuf::new(),
SessionOptions {
disable_dht: !argument.enable_dht,
disable_upload: !argument.enable_upload,
persistence: None,
trackers: trackers.clone(),
..SessionOptions::default()
},
)
.await?;
// collect info-hashes from API
for source in &argument.infohash_file {
if is_debug_i {
debug::info(format!("Handle info-hash source `{source}`..."));
}
// collect info-hashes from each API channel
for source in &arg.infohash_file {
debug.info(&format!("Handle info-hash source `{source}`..."));
// aquatic server may update the stats at this moment,
// handle this state manually
match api::infohashes(source) {
Ok(infohashes) => {
total += infohashes.len();
for i in infohashes {
if torrent_storage.as_ref().is_some_and(|s| !s.exists(&i)) {
if is_debug_i {
debug::info(format!("Resolve `{i}`..."));
}
match session
.add_torrent(
librqbit::AddTorrent::from_url(format!(
"magnet:?xt=urn:btih:{i}"
)),
Some(librqbit::AddTorrentOptions {
disable_trackers: trackers.is_empty(),
initial_peers: if peers.is_empty() {
None
} else {
Some(peers.clone())
},
list_only: true,
..Default::default()
}),
)
.await?
{
librqbit::AddTorrentResponse::ListOnly(r) => {
if let Some(ref s) = torrent_storage {
let p = s.save(&i, &r.torrent_bytes)?;
if is_debug_i {
debug::info(format!(
"Add new torrent file `{}`",
p.to_string_lossy()
));
debug.info(&format!("Index `{i}`..."));
match session
.add_torrent(
AddTorrent::from_url(format!("magnet:?xt=urn:btih:{i}")),
Some(AddTorrentOptions {
overwrite: true,
disable_trackers: trackers.is_empty(),
initial_peers: if peers.is_empty() {
None
} else {
Some(peers.clone())
},
// preload nothing, but listing when regex pattern argument is given
list_only: arg.preload_regex.is_none(),
// this option allows rqbit manager to preload some or any files match pattern
// * useful to build index with multimedia files, like images for audio albums
output_folder: storage.output_folder(&i).ok(),
// applies preload some files to the destination directory (above)
only_files_regex: arg.preload_regex.clone(),
..Default::default()
}),
)
.await
{
Ok(r) => match r {
AddTorrentResponse::AlreadyManaged(_, t)
| AddTorrentResponse::Added(_, t) => {
if arg.save_torrents {
t.with_metadata(|m| {
save_torrent_file(
&storage,
&debug,
&i,
&m.torrent_bytes,
)
})?;
}
/*tokio::spawn({
let t = t.clone();
let d = Duration::from_secs(5);
async move {
loop {
let s = t.stats();
if s.finished {
break;
}
debug.info(&format!("{s}..."));
tokio::time::sleep(d).await;
}
}
});*/
// @TODO t.wait_until_completed().await?;
}
AddTorrentResponse::ListOnly(r) => {
if arg.save_torrents {
save_torrent_file(&storage, &debug, &i, &r.torrent_bytes)
}
// @TODO
// use `r.info` for Memory, SQLite, Manticore and other alternative storage type
}
_ => panic!(),
}
},
Err(e) => debug.info(&format!("Torrent handle skipped: `{e}`")),
}
}
}
Err(ref e) => {
if is_debug_e {
debug::error(e)
}
}
Err(e) => debug.error(&e.to_string()),
}
}
session.stop().await;
if is_debug_i {
debug::info(format!(
"Index of {total} hashes completed, await {} seconds to continue...",
argument.sleep,
));
}
std::thread::sleep(std::time::Duration::from_secs(argument.sleep));
debug.info(&format!(
"Index of {total} hashes completed, await {} seconds to continue...",
arg.sleep,
));
std::thread::sleep(Duration::from_secs(arg.sleep));
}
}
fn save_torrent_file(s: &Storage, d: &Debug, i: &str, b: &[u8]) {
if s.torrent_exists(i) {
d.info(&format!("Torrent file `{i}` already exists, skip"))
} else {
match s.save_torrent(i, b) {
Ok(r) => d.info(&format!("Add torrent file `{}`", r.to_string_lossy())),
Err(e) => d.error(&e.to_string()),
}
}
}

21
src/peers.rs Normal file
View file

@ -0,0 +1,21 @@
use std::{net::SocketAddr, str::FromStr};
pub struct Peers(Vec<SocketAddr>);
impl Peers {
pub fn init(peers: &Vec<String>) -> anyhow::Result<Self> {
let mut p = Vec::with_capacity(peers.len());
for peer in peers {
p.push(SocketAddr::from_str(peer)?);
}
Ok(Self(p))
}
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
pub fn clone(&self) -> Vec<SocketAddr> {
self.0.clone()
}
}

61
src/storage.rs Normal file
View file

@ -0,0 +1,61 @@
use anyhow::{Result, bail};
use std::{fs, io::Write, path::PathBuf, str::FromStr};
pub struct Storage(PathBuf);
impl Storage {
pub fn init(storage: &str, clear: bool) -> Result<Self> {
let p = PathBuf::from_str(storage)?;
if let Ok(t) = fs::metadata(&p) {
if t.is_file() {
bail!("Storage destination is not directory!");
}
if t.is_dir() && clear {
for i in fs::read_dir(&p)? {
let r = i?.path();
if r.is_dir() {
fs::remove_dir_all(&r)?;
} else {
fs::remove_file(&r)?;
}
}
}
}
fs::create_dir_all(&p)?;
Ok(Self(p))
}
pub fn torrent_exists(&self, infohash: &str) -> bool {
fs::metadata(self.torrent(infohash))
.is_ok_and(|p| p.is_file() || p.is_dir() || p.is_symlink())
}
pub fn save_torrent(&self, infohash: &str, bytes: &[u8]) -> Result<PathBuf> {
let p = self.torrent(infohash);
let mut f = fs::File::create(&p)?;
f.write_all(bytes)?;
Ok(p)
}
pub fn output_folder(&self, infohash: &str) -> Result<String> {
let mut p = PathBuf::new();
p.push(&self.0);
p.push(infohash);
if p.is_file() {
bail!("File destination is not directory!");
}
fs::create_dir_all(&p)?;
Ok(p.to_string_lossy().to_string())
}
pub fn path(&self) -> PathBuf {
self.0.clone()
}
fn torrent(&self, infohash: &str) -> PathBuf {
let mut p = PathBuf::new();
p.push(&self.0);
p.push(format!("{infohash}.torrent"));
p
}
}

20
src/trackers.rs Normal file
View file

@ -0,0 +1,20 @@
use std::{collections::HashSet, str::FromStr};
use url::Url;
pub struct Trackers(HashSet<Url>);
impl Trackers {
pub fn init(trackers: &Vec<String>) -> anyhow::Result<Self> {
let mut t = HashSet::with_capacity(trackers.len());
for tracker in trackers {
t.insert(Url::from_str(tracker)?);
}
Ok(Self(t))
}
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
pub fn clone(&self) -> HashSet<Url> {
self.0.clone()
}
}