domainscan 0.2.0

use std::net::IpAddr;
use std::time::{Duration, Instant};
use std::sync::{Mutex, Arc};
use std::sync::mpsc::{channel ,Sender, Receiver};
use futures::{stream, StreamExt};
use tokio::time::{timeout};
use trust_dns_resolver::AsyncResolver;
use trust_dns_resolver::config::*;
use crate::model::Domain;
use crate::result::{ScanStatus, DomainScanResult};

#[cfg(feature = "passive")]
use reqwest::Url;
#[cfg(feature = "passive")]
use crate::model::CertEntry;
#[cfg(feature = "passive")]
use crate::config::{URL_CRT, DEFAULT_USER_AGENT};

/// Structure for domain scan  
/// 
/// Should be constructed using DomainScanner::new
#[derive(Clone)]
pub struct DomainScanner {
    /// Base Domain Name of scan target.  
    base_domain: String,
    /// Word-list of name
    word_list: Vec<String>,
    /// Timeout setting of domain scan.  
    timeout: Duration,
    /// Result of domain scan.  
    scan_result: DomainScanResult,
    /// Sender for progress messaging
    tx: Arc<Mutex<Sender<String>>>,
    /// Receiver for progress messaging
    rx: Arc<Mutex<Receiver<String>>>,
    /// Run passive scan
    passive: bool,
}

impl DomainScanner {
    /// Construct new UriScanner  
    pub fn new() -> Result<DomainScanner, String> {
        let (tx, rx) = channel();
        let domain_scanner = DomainScanner {
            base_domain: String::new(),
            word_list: vec![],
            timeout: Duration::from_millis(30000),
            scan_result: DomainScanResult::new(),
            tx: Arc::new(Mutex::new(tx)),
            rx: Arc::new(Mutex::new(rx)),
            passive: false,
        };
        Ok(domain_scanner)
    }
    /// Set base Domain of scan target.  
    pub fn set_base_domain(&mut self, base_domain: String) {
        self.base_domain = base_domain;
    }
    /// Add word to word-list
    pub fn add_word(&mut self, word: String) {
        self.word_list.push(word);
    }
    /// Set word-list 
    pub fn set_word_list(&mut self, word_list: Vec<&str>) {
        self.word_list.clear();
        for word in word_list {
            self.word_list.push(word.to_string())
        }
    }
    /// Set scan timeout  
    pub fn set_timeout(&mut self, timeout: Duration){
        self.timeout = timeout;
    }
    /// Set active/passive scan (default is active)
    pub fn set_passive(&mut self, passive: bool) {
        self.passive = passive;
    }
    async fn scan_domain(&self) -> Result<Vec<Domain>, ()> {
        if self.passive {
            #[cfg(feature = "passive")]
            match timeout(self.timeout, scan_subdomain_passive(self.base_domain.clone(), &self.tx)).await {
                Ok(domains) => {
                    return Ok(domains);
                },
                Err(_) => {
                    return Err(());
                },
            }
            #[cfg(not(feature = "passive"))]
            return Err(());
        }else{
            match timeout(self.timeout, scan_subdomain(self.base_domain.clone(), self.word_list.clone(), &self.tx)).await {
                Ok(domains) => {
                    return Ok(domains);
                },
                Err(_) => {
                    return Err(());
                },
            }
        }
    }
    /// Run scan with current settings. 
    /// 
    /// Results are stored in DomainScanner::scan_result
    pub async fn run_scan(&mut self){
        if self.passive && cfg!(not(feature="passive")) {
            self.scan_result.scan_status = ScanStatus::Error;
            return; 
        }
        let start_time = Instant::now();
        let res = self.scan_domain().await;
        match res {
            Ok(domains) => {
                self.scan_result.domains = domains;
                self.scan_result.scan_status = ScanStatus::Done;
            },
            Err(_) => {
                self.scan_result.scan_status = ScanStatus::Timeout;
            },
        }
        self.scan_result.scan_time = Instant::now().duration_since(start_time);
    }
    /// Return scan result.
    pub fn get_result(&mut self) -> DomainScanResult{
        return self.scan_result.clone();
    }
    /// Run scan and return result
    pub async fn scan(&mut self) -> DomainScanResult {
        self.run_scan().await;
        self.scan_result.clone()
    }
    /// Get progress receiver
    pub fn get_progress_receiver(&self) -> Arc<Mutex<Receiver<String>>> {
        self.rx.clone()
    }
}

async fn resolve_domain(domain: String) -> Vec<IpAddr> {
    let mut ips: Vec<IpAddr> = vec![];
    let resolver = AsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default()).unwrap();
    match resolver.lookup_ip(domain).await {
        Ok(lip) => {
            for ip in lip.iter() {
                ips.push(ip);
            }
        },
        Err(_) => {},
    }
    ips
}

#[cfg(feature = "passive")]
fn extract_domain(target: String) -> String {
    let mut domain_name: String = target;
    match domain_name.strip_prefix("*.") {
        Some(d) => {
            domain_name = d.to_string();
        },
        None => {},
    }
    domain_name
}

#[cfg(feature = "passive")]
fn is_subdomain(domain: String, apex_domain: String) -> bool {
    domain.contains(&apex_domain) && domain.ends_with(&apex_domain) && domain.len() > apex_domain.len() 
}

async fn scan_subdomain(base_domain: String, word_list: Vec<String>, ptx: &Arc<Mutex<Sender<String>>>) -> Vec<Domain> {
    let mut result: Vec<Domain> = vec![];
    let scan_results: Arc<Mutex<Vec<Domain>>> = Arc::new(Mutex::new(vec![]));
    let mut target_domains: Vec<String> = vec![];
    for word in word_list {
        target_domains.push(format!("{}.{}", word, base_domain));
    }
    let results = stream::iter(target_domains).map(|domain| {
        async move {
            let ips: Vec<IpAddr> = resolve_domain(domain.clone()).await;
            let d = Domain {
                domain_name: domain.clone(),
                ips: ips,
            };
            match ptx.lock() {
                Ok(lr) => {
                    match lr.send(domain) {
                        Ok(_) => {},
                        Err(_) => {},
                    }
                },
                Err(_) => {},
            }
            d
        }
    }).buffer_unordered(100);
    results.for_each(|domain| async {
        if domain.ips.len() > 0 {
            scan_results.lock().unwrap().push(domain);
        }
    }).await;
    for domain in scan_results.lock().unwrap().iter() {
        result.push(domain.to_owned());
    }
    result
}

#[cfg(feature = "passive")]
async fn scan_subdomain_passive(base_domain: String, ptx: &Arc<Mutex<Sender<String>>>) -> Vec<Domain>  {
    let mut result: Vec<Domain> = vec![];
    let scan_results: Arc<Mutex<Vec<Domain>>> = Arc::new(Mutex::new(vec![]));
    let mut certs: Vec<CertEntry> = vec![];
    //"https://crt.sh/?dNSName=example.com&output=json"
    let url = match Url::parse_with_params(URL_CRT, &[("dNSName", base_domain.clone().as_str()), ("output", "json")]){
        Ok(url) => url,
        Err(e) => {
            println!("{}",e);
            return result;
        },
    };
    let client = reqwest::Client::builder().timeout(Duration::from_secs(60)).build().expect("failed to build HTTP reqest client");
    let res = client.get(url).header(reqwest::header::USER_AGENT, DEFAULT_USER_AGENT).send().await;
    match res {
        Ok(r) => {
            if r.status().is_success() {
                match r.text().await {
                    Ok(res_text) => {
                        let certs_json:serde_json::Value = serde_json::from_str(res_text.as_str()).unwrap();
                        if certs_json.is_array() {
                            let cert_array = certs_json.as_array().unwrap();
                            for cert in cert_array {
                                match serde_json::to_string(cert) {
                                    Ok(cert) => {
                                        let cert:CertEntry = match serde_json::from_str(cert.as_str()){
                                            Ok(cert) => cert,
                                            Err(_) => continue,
                                        };
                                        certs.push(cert);
                                    },
                                    Err(_) => {},
                                }
                            }
                        }
                    },
                    Err(_) => {},
                };
            }
        },
        Err(_) => {},
    }
    let mut target_domains: Vec<String> = vec![];
    for cert in certs {
        let domain_name: String = extract_domain(cert.common_name);
        if is_subdomain(domain_name.clone(), base_domain.clone()) && !target_domains.contains(&domain_name) {
            target_domains.push(domain_name);
        }
        let name_values: Vec<&str> = cert.name_value.trim().split("\n").collect();
        for value in name_values {
            let name: String = extract_domain(value.to_string());
            if is_subdomain(name.clone(), base_domain.clone()) && !target_domains.contains(&name) {
                target_domains.push(name);
            }
        }
    }
    let results = stream::iter(target_domains).map(|domain| {
        async move {
            let ips: Vec<IpAddr> = resolve_domain(domain.clone()).await;
            let d = Domain {
                domain_name: domain.clone(),
                ips: ips,
            };
            match ptx.lock() {
                Ok(lr) => {
                    match lr.send(domain) {
                        Ok(_) => {},
                        Err(_) => {},
                    }
                },
                Err(_) => {},
            }
            d
        }
    }).buffer_unordered(100);
    results.for_each(|domain| async {
        if domain.ips.len() > 0 {
            scan_results.lock().unwrap().push(domain);
        }
    }).await;
    for domain in scan_results.lock().unwrap().iter() {
        result.push(domain.to_owned());
    }
    result
}