1
0
Fork 0
forked from zesterer/babble
babble/src/main.rs

144 lines
4.1 KiB
Rust

mod generator;
use axum::{Router, extract::Path, response::Html, routing::get};
use axum_server::{bind_rustls, tls_rustls::RustlsConfig};
use clap::Parser;
use hashbrown::HashMap;
use itertools::Itertools;
use rand::{Rng as _, prelude::*, random_range};
use rand_chacha::{ChaCha8Rng, rand_core::SeedableRng};
use std::{
borrow::Cow,
path::PathBuf,
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
pub type Rng = ChaCha8Rng;
#[derive(Parser)]
pub struct Args {
/// Socket to bind to, defaults to 0.0.0.0:3000
#[arg(long)]
sock: Option<String>,
#[arg(long)]
/// Path of the certificate .pem
cert: Option<PathBuf>,
/// Path of the private key .pem
#[arg(long)]
key: Option<PathBuf>,
}
#[tokio::main]
async fn main() {
let args = Args::parse();
let generators: Vec<Arc<dyn generator::Generator>> = vec![
Arc::new(generator::Markov::new(include_str!("../wap.txt"))),
Arc::new(generator::Ast::new()),
];
let counter = Arc::new(AtomicU64::new(0));
let app = {
let counter = counter.clone();
Router::new().route(
"/{id}",
get(|Path(id): Path<String>| async move {
tokio::time::sleep(Duration::from_millis(random_range(200..1000))).await;
counter.fetch_add(1, Ordering::Relaxed);
let mut seed = [0; 32];
for (i, b) in id
.bytes()
.chain(core::iter::repeat(0))
.take(seed.len())
.enumerate()
{
seed[i] = b;
}
let mut rng = Rng::from_seed(seed);
let generator = generators.choose(&mut rng).unwrap();
let title = generator
.word_stream(rng.random_range(2..10), &mut rng.clone())
.join(" ");
let content = generator
.word_stream(rng.random_range(50..5_000), &mut rng.clone())
.fold(String::new(), |mut content, word| {
if rng.random_bool(0.05) {
let url = generator.word_stream(3, &mut rng.clone()).join("-");
content += &format!(" <a href=\"{}\">{}</a>", url, word);
} else if rng.random_bool(0.01) {
content += ".<br>";
} else {
content += " ";
content += &word
}
content
});
Html(format!(
"<!DOCTYPE html>
<html>
<head>
<title>{title}</title>
</head>
<body>
<h1>{title}</h1>
<p>{content}</p>
</body>
</html>"
))
}),
)
};
let mut interval = tokio::time::interval(Duration::from_secs(20));
tokio::spawn(async move {
let mut last = 0;
loop {
interval.tick().await;
let count = counter.load(Ordering::Relaxed);
if count != last {
last = count;
println!(
"{} Served bollocks to {} clients!",
chrono::offset::Local::now(),
count,
);
}
}
});
println!("Starting...");
let sock = args
.sock
.as_deref()
.unwrap_or("0.0.0.0:3000")
.parse()
.unwrap();
if let (Some(cert), Some(key)) = (args.cert, args.key) {
println!("Enabling TLS...");
let config = RustlsConfig::from_pem_file(cert, key).await.unwrap();
bind_rustls(sock, config)
.serve(app.into_make_service())
.await
.unwrap();
} else {
println!("WARNING: TLS disabled.");
axum_server::bind(sock)
.serve(app.into_make_service())
.await
.unwrap()
}
}