use super::error;
use crate::models::host::Host;
use crate::resolvers;
use serde::{Deserialize, Serialize};
use std::{
net::IpAddr,
sync::{Arc, Mutex},
};
use trust_dns_resolver::config::ResolverConfig;
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Scan {
pub hosts: Vec<Host>,
pub domain: String,
pub ticks: i64,
}
impl Scan {
pub fn new(domain: String) -> Self {
let domain = format!("{}.", domain);
Scan {
hosts: vec![],
domain,
ticks: 0,
}
}
pub fn new_arc_mutex(domain: String) -> Arc<Mutex<Self>> {
Arc::new(Mutex::new(Scan::new(domain)))
}
}
impl Scan {
pub fn inc_tick(&mut self, i: i64) {
self.ticks += i;
}
pub fn is_tick_available(&self, len: &usize) -> bool {
return &usize::try_from(self.ticks).unwrap() < len;
}
}
impl Scan {
pub fn contains_host(&self, name: &String) -> bool {
let names = self
.hosts
.iter()
.map(|h| h.name.clone())
.collect::<Vec<String>>();
names.contains(&name)
}
pub fn add_host(&mut self, name: String) {
let host = Host::new(name);
self.hosts.push(host);
}
pub fn add_ip_for_host(&mut self, name: &String, ip: IpAddr) {
for host in &mut self.hosts {
if &host.name == name {
if !host.ips.contains(&ip) {
host.ips.push(ip);
}
}
}
}
pub fn get_host_by_name(&self, name: &String) -> &Host {
return self
.hosts
.iter()
.filter(|h| &h.name == name)
.collect::<Vec<&Host>>()[0];
}
pub fn host_contains_ip(&self, name: &String, ip: &IpAddr) -> bool {
let host: &Host = self.get_host_by_name(name);
if host.ips.contains(ip) {
return true;
}
return false;
}
}
impl Scan {
pub fn get_resolver_config(name: &String) -> ResolverConfig {
let resolver_config: ResolverConfig;
if name == resolvers::CLOUDFLARE {
resolver_config = ResolverConfig::cloudflare();
} else if name == resolvers::QUAD {
resolver_config = ResolverConfig::quad9();
} else {
resolver_config = ResolverConfig::google();
}
resolver_config
}
}
impl Scan {
pub fn to_json(&self) -> Result<String, error::Error> {
match serde_json::to_string(self) {
Ok(res) => Ok(res),
Err(e) => Err(error::Error::from(e)),
}
}
}
impl Scan {
pub fn to_vec(&self) -> Vec<[String; 2]> {
let mut string_vec = vec![];
for host in &self.hosts {
for ip in &host.ips {
string_vec.push([host.name.clone(), ip.to_string()])
}
}
return string_vec;
}
}
impl Scan {
pub fn to_string_with_sep(&self, sep: &str) -> String {
self.to_vec()
.iter()
.map(|r| r.join(sep))
.collect::<Vec<String>>()
.join("\n")
}
}
impl ToString for Scan {
fn to_string(&self) -> String {
self.to_string_with_sep(" ")
}
}
impl Scan {
pub fn to_csv(&self) -> String {
return vec!["Subdomain,Ip".to_string(), self.to_string_with_sep(",")].join("\n");
}
}