use std::{future::Future, pin::Pin};
use actix_web::HttpRequest;
use crate::{leaked, models::{app::MiddlewareNext, response::{Response, Status}}, util::req::extract_real_ip};
use colored::Colorize;
#[derive(Clone)]
pub enum RockUser {
IP(&'static str),
Range(&'static str),
UserAgent(&'static str),
UserAgentPortion(&'static str),
}
#[derive(Clone)]
pub struct RockArgs {
pub total_blacklist: Vec<RockUser>,
pub blacklist_paths: Vec<(RockUser, Box<[&'static str]>)>,
pub whitelist_paths: Vec<(RockUser, Box<[&'static str]>)>,
pub total_redirects: Vec<(RockUser, &'static str)>,
}
impl RockArgs {
pub fn new() -> Self {
Self {
total_blacklist: Vec::new(),
blacklist_paths: Vec::new(),
whitelist_paths: Vec::new(),
total_redirects: Vec::new(),
}
}
pub fn add_blacklist(&mut self, user: RockUser) {
self.total_blacklist.push(user);
}
pub fn add_blacklist_path(&mut self, user: RockUser, paths: Box<[&'static str]>) {
self.blacklist_paths.push((user, paths));
}
pub fn add_whitelist_path(&mut self, user: RockUser, paths: Box<[&'static str]>) {
self.whitelist_paths.push((user, paths));
}
pub fn add_redirect(&mut self, user: RockUser, path: &'static str) {
self.total_redirects.push((user, path));
}
}
pub fn load_blocklist(path: &'static str) -> RockArgs {
let mut args = RockArgs::new();
let lines = std::fs::read_to_string(path).unwrap();
let modified_lines = lines.replace("\r", "");
let lines_clone: Vec<String> = modified_lines.lines().map(|line| line.to_string()).collect();
for line in lines_clone {
if str::starts_with(&line, "#") {
continue;
}
if str::trim(&line).is_empty() {
continue;
}
let parts = leaked(line).split_whitespace().collect::<Vec<&str>>();
if parts.len() == 1 {
args.add_blacklist(RockUser::IP(parts[0]));
} else {
args.add_blacklist(RockUser::IP(parts[1]));
}
}
args
}
pub fn rock(args: RockArgs) -> impl Fn(&HttpRequest) -> Pin<Box<dyn Future<Output = Result<Response, MiddlewareNext>>>> {
let args = args.clone();
let mut robots_txt: Vec<String> = Vec::new();
let mut applied_uas: Vec<String> = Vec::new();
for user in args.total_blacklist.iter() {
match user {
RockUser::UserAgent(user_agent) => {
let ua = "User-Agent: ".to_string() + user_agent;
robots_txt.push(ua);
robots_txt.push("Disallow: /".to_string());
applied_uas.push(user_agent.to_string());
},
RockUser::UserAgentPortion(user_agent) => {
let ua = "User-Agent: ".to_string() + user_agent;
robots_txt.push(ua);
robots_txt.push("Disallow: /".to_string());
applied_uas.push(user_agent.to_string());
},
_ => {},
}
}
for (user, paths) in args.blacklist_paths.iter() {
match user {
RockUser::UserAgent(user_agent) => {
if applied_uas.iter().any(|ua| ua == user_agent) {
continue;
}
let ua = "User-Agent: ".to_string() + user_agent;
robots_txt.push(ua.clone());
for path in paths.iter() {
let disallow = "Disallow: ".to_string() + path;
robots_txt.push(disallow);
}
},
RockUser::UserAgentPortion(user_agent) => {
if applied_uas.iter().any(|ua| ua == user_agent) {
continue;
}
let ua = "User-Agent: ".to_string() + user_agent;
robots_txt.push(ua.clone());
for path in paths.iter() {
let disallow = "Disallow: ".to_string() + path;
robots_txt.push(disallow);
}
},
_ => {},
}
}
let robots_txt = robots_txt.join("\n");
return move |req: &HttpRequest| {
let args_clone = args.clone();
let req_clone = req.clone();
let txt_clone = robots_txt.clone();
let ip = extract_real_ip(req);
if ip.is_none() {
println!("[{}] {} -> Failed to extract IP", "ROCK".bright_black(), "ERROR".red());
return Box::pin(async move {
Ok(Response::new(Status::from_u16(500)))
});
}
let ip = ip.unwrap();
return Box::pin(async move {
let args = args_clone.clone();
let ip = &ip.clone();
let req = req_clone.clone();
let txt = txt_clone.clone();
if req.path() == "/robots.txt" {
let mut resp = Response::new(Status::from_u16(200));
resp.body_text(txt);
return Ok(resp);
}
for (user, paths) in args.blacklist_paths.iter() {
match user {
RockUser::IP(user_ip) => {
if ip == user_ip && paths.iter().any(|p| str::starts_with(req_clone.path(), p)) {
println!("[{}] {} {} -> Failed IP {}", "ROCK".bright_black(), "Blocked".red(), ip.bright_blue(), user_ip.bright_blue());
return Ok(Response::new(Status::from_u16(403)));
}
},
RockUser::Range(range) => {
if str::starts_with(ip, range) && paths.iter().any(|p| str::starts_with(req_clone.path(), p)) {
println!("[{}] {} {} -> Failed RANGE {}", "ROCK".bright_black(), "Blocked".red(), ip.bright_blue(), range.bright_blue());
return Ok(Response::new(Status::from_u16(403)));
}
},
RockUser::UserAgent(user_agent) => {
if &req.headers().get("User-Agent").unwrap().to_str().unwrap() == user_agent && paths.iter().any(|p| str::starts_with(req_clone.path(), p)) {
println!("[{}] {} {} -> Failed UA {}", "ROCK".bright_black(), "Blocked".red(), ip.bright_blue(), user_agent.bright_magenta());
return Ok(Response::new(Status::from_u16(403)));
}
},
RockUser::UserAgentPortion(user_agent) => {
if req.headers().get("User-Agent").unwrap().to_str().unwrap().contains(user_agent) && paths.iter().any(|p| str::starts_with(req_clone.path(), p)) {
println!("[{}] {} {} -> Failed UAP {}", "ROCK".bright_black(), "Blocked".red(), ip.bright_blue(), user_agent.bright_magenta());
return Ok(Response::new(Status::from_u16(403)));
}
},
}
}
for (user, paths) in args.whitelist_paths.iter() {
match user {
RockUser::IP(user_ip) => {
if ip == user_ip && paths.iter().any(|p| str::starts_with(req_clone.path(), p)) {
return Ok(Response::new(Status::from_u16(200)));
}
},
RockUser::Range(range) => {
if str::starts_with(ip, range) && paths.iter().any(|p| str::starts_with(req_clone.path(), p)) {
return Ok(Response::new(Status::from_u16(200)));
}
},
RockUser::UserAgent(user_agent) => {
if &req.headers().get("User-Agent").unwrap().to_str().unwrap() == user_agent && paths.iter().any(|p| str::starts_with(req_clone.path(), p)) {
return Ok(Response::new(Status::from_u16(200)));
}
},
RockUser::UserAgentPortion(user_agent) => {
if req.headers().get("User-Agent").unwrap().to_str().unwrap().contains(user_agent) && paths.iter().any(|p| str::starts_with(req_clone.path(), p)) {
return Ok(Response::new(Status::from_u16(200)));
}
},
}
}
for user in args.total_blacklist.iter() {
match user {
RockUser::IP(user_ip) => {
if &ip == user_ip {
println!("[{}] {} {} -> Failed IP {}", "ROCK".bright_black(), "TBlocked".red(), ip.bright_blue(), user_ip.bright_blue());
return Ok(Response::new(Status::from_u16(403)));
}
},
RockUser::Range(range) => {
if str::starts_with(&ip, range) {
println!("[{}] {} {} -> Failed RANGE {}", "ROCK".bright_black(), "TBlocked".red(), ip.bright_blue(), range.bright_blue());
return Ok(Response::new(Status::from_u16(403)));
}
},
RockUser::UserAgent(user_agent) => {
if &req.headers().get("User-Agent").unwrap().to_str().unwrap() == user_agent {
println!("[{}] {} {} -> Failed UA {}", "ROCK".bright_black(), "TBlocked".red(), ip.bright_blue(), user_agent.bright_magenta());
return Ok(Response::new(Status::from_u16(403)));
}
},
RockUser::UserAgentPortion(user_agent) => {
if req.headers().get("User-Agent").unwrap().to_str().unwrap().contains(user_agent) {
println!("[{}] {} {} -> Failed UAP {}", "ROCK".bright_black(), "TBlocked".red(), ip.bright_blue(), user_agent.bright_magenta());
return Ok(Response::new(Status::from_u16(403)));
}
},
}
}
for (user, path) in args.total_redirects.iter() {
if req.path() == *path {
return Err(MiddlewareNext::Next);
}
match user {
RockUser::IP(user_ip) => {
if &ip == user_ip {
println!("[{}] {} {} -> Redirected IP {}", "ROCK".bright_black(), "Redirected".yellow(), ip.bright_blue(), path.bright_yellow());
let mut resp = Response::new(Status::from_u16(302));
resp.redirect(path.to_string());
return Ok(resp);
}
},
RockUser::Range(range) => {
if str::starts_with(&ip, range) {
println!("[{}] {} {} -> Redirected RANGE {}", "ROCK".bright_black(), "Redirected".yellow(), ip.bright_blue(), path.bright_yellow());
let mut resp = Response::new(Status::from_u16(302));
resp.redirect(path.to_string());
return Ok(resp);
}
},
RockUser::UserAgent(user_agent) => {
if &req.headers().get("User-Agent").unwrap().to_str().unwrap() == user_agent {
println!("[{}] {} {} -> Redirected UA {}", "ROCK".bright_black(), "Redirected".yellow(), ip.bright_blue(), user_agent.bright_magenta());
let mut resp = Response::new(Status::from_u16(302));
resp.redirect(path.to_string());
return Ok(resp);
}
},
RockUser::UserAgentPortion(user_agent) => {
if req.headers().get("User-Agent").unwrap().to_str().unwrap().contains(user_agent) {
println!("[{}] {} {} -> Redirected UAP {}", "ROCK".bright_black(), "Redirected".yellow(), ip.bright_blue(), user_agent.bright_magenta());
let mut resp = Response::new(Status::from_u16(302));
resp.redirect(path.to_string());
return Ok(resp);
}
},
}
}
Err(MiddlewareNext::Next)
});
};
}