ascertain/src/report.rs

use std::borrow::Cow;
use std::collections::HashSet;
use std::future::Future;
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
use std::path::PathBuf;
use std::sync::Arc;

use base64::Engine;
use base64::prelude::BASE64_STANDARD;
use chrono::{Local, Utc};
use openssl::asn1::{Asn1Time, Asn1TimeRef};
use openssl::hash::MessageDigest;
use openssl::pkey::{Id, PKeyRef, Public};
use openssl::stack::StackRef;
use openssl::x509::{GeneralNameRef, X509, X509NameEntryRef, X509NameRef, X509Ref};
use serde::{Deserialize, Serialize};
use serde_with::{base64::Base64, serde_as};
use thiserror::Error;
use tokio::io::{AsyncWriteExt, BufWriter};
use tokio::sync::{mpsc, RwLock};
use tokio::sync::mpsc::Sender;
use tracing::{error, info, warn};

use crate::config::OutputFormat;

#[derive(Error, Debug, Serialize, Copy, Clone)]
pub enum  ReportError {
    #[error("Connection timed out")]
    ConnectionTimeout,
    #[error("Connection refused")]
    ConnectionRefused,
    #[error("Handshake timed out")]
    HandshakeTimeout,
    #[error("TLS Protocol error: probably not a TLS server")]
    ProtocolError,
}

#[derive(Serialize, Debug, Clone)]
#[serde(rename="lowercase", tag="status", content="report")]
pub enum ReportPayload {
    Success{
        certificate: CertInfo,
        // TODO: define fields for SSL implementation
    },
    Error{msg: ReportError},
}

#[derive(Serialize, Debug, Clone)]
pub struct ProbeReport {
    pub host: SocketAddr,
    pub scan_date: chrono::DateTime<Local>,
    #[serde(flatten)]
    pub result: ReportPayload
}


#[serde_as]
#[derive(Serialize, Debug, Clone)]
pub struct CertInfo {
    #[serde_as(as="Base64")]
    pub cert_digest: Vec<u8>,
    #[serde_as(as="Base64")]
    pub issuer_subject_der: Vec<u8>,
    pub issuer_subject: Vec<String>,
    #[serde_as(as="Base64")]
    pub certificate_der: Vec<u8>,
    pub subject: Vec<String>,
    #[serde_as(as="Base64")]
    pub subject_der: Vec<u8>,
    pub san: Vec<String>,
    pub not_before: chrono::DateTime<Utc>,
    pub not_after: chrono::DateTime<Utc>,
    pub key_type: String,
    pub signature_type: String,
    #[serde(with="hex")]
    pub authority_key_id: Vec<u8>,
    #[serde(with="hex")]
    pub subject_key_id: Vec<u8>,
    pub verification_state: VerificationState,
}

#[derive(Serialize, Debug, Copy, Clone, Eq, PartialEq)]
pub enum VerificationState {
    Unknown,
    Valid,
    Invalid,
}
fn asn1time_to_datetime(date: &Asn1TimeRef) -> anyhow::Result<chrono::DateTime<Utc>> {
    let res = Asn1Time::from_unix(0).unwrap().diff(date)?;
    let timestamp = res.days as i64 * 86400 + res.secs as i64;
    chrono::DateTime::from_timestamp(timestamp, 0)
        .ok_or(anyhow::anyhow!("Constructing timestamp failed"))
}

fn describe_key(key: &PKeyRef<Public>) -> String {
    match key.id() {
        Id::RSA => format!("RSA-{}", key.bits()),
        Id::RSA_PSS => format!("RSA-PSS-{}", key.bits()),
        Id::DSA => format!("DSA-{}", key.bits()),
        Id::EC => format!("EC-P{}", key.bits()),
        Id::ED25519 => "Ed25519".to_owned(),
        Id::ED448 => "Ed448".to_owned(),
        id => format!("UNKNOWN-{}", id.as_raw()),
    }
}

impl CertInfo {
    pub fn extract(data: &X509Ref) -> anyhow::Result<Self> {
        let md = MessageDigest::sha256();
        let cert_digest = data.digest(md)?.to_vec();
        let issuer_subject = data.issuer_name();


        Ok(CertInfo {
            cert_digest,
            issuer_subject_der: issuer_subject.to_der()?,
            issuer_subject: issuer_subject.entries().map(format_x509_name_entry).collect(),
            certificate_der: data.to_der()?,
            subject: data.subject_name().entries().map(format_x509_name_entry).collect(),
            subject_der: data.subject_name().to_der()?,
            san: data.subject_alt_names()
                .map(|stack| stack.iter().map(format_general_name).collect())
                .unwrap_or(Vec::new()),
            not_before: asn1time_to_datetime(data.not_before())?,
            not_after: asn1time_to_datetime(data.not_after())?,
            key_type: describe_key(data.public_key()?.as_ref()),
            signature_type: data.signature_algorithm().object().nid().short_name()?.to_owned(),
            authority_key_id: data.authority_key_id().map_or(Vec::new(), |id| id.as_slice().to_vec()),
            subject_key_id: data.subject_key_id().map_or(Vec::new(), |id| id.as_slice().to_vec()),
            verification_state: VerificationState::Unknown,
        })
    }
}

#[derive(Deserialize)]
pub struct JsonConfig {
    /// File which receives discovered certificates.
    output_file: PathBuf,

    /// File which receives discovered issuer certificates. Optional; if not included, do not store issuers.
    #[serde(default)]
    issuer_file: Option<PathBuf>,

    /// Enable an outer container so that the output file contains a single JSON list.
    /// If disabled (default), the file is in json-lines format.
    #[serde(default)]
    container: bool,
}

#[allow(unused)]
#[derive(Clone)]
pub struct Reporter {
    issuer_chan: Sender<X509>,
    report_chan: Sender<ProbeReport>,
    digests: Arc<RwLock<HashSet<Vec<u8>>>>,
    collect_digests: bool,
}

#[derive(Error, Debug)]
pub enum ReportingError {
    #[error("Report formatter terminated")]
    ReportFormatterFailed,
}

fn format_x509_name_entry(entry: &X509NameEntryRef) -> String {
    let name = entry.object().nid().short_name()
        .map(Cow::Borrowed)
        .unwrap_or_else(|_| Cow::Owned(format!("{:?}", entry.object())));
    let value = entry.data().as_utf8().map(|data| data.to_string()).unwrap_or_else(|_|
        BASE64_STANDARD.encode(entry.data().as_slice())
    );
    format!("{name}={value}")
}
fn format_x509_name(name: &X509NameRef) -> String {
    let mut result = "".to_owned();
    for entry in name.entries() {
        if !result.is_empty() {
            result.push_str(", ")
        }
        result.push_str(format_x509_name_entry(entry).as_ref());
    }
    result
}

fn format_general_name(name: &GeneralNameRef) -> String {
    // TODO: there's other types that aren't supported by the safe wrapper.
    if let Some(name) = name.email() {
        return format!("EMAIL:{name}")
    } else if let Some(name) = name.dnsname() {
        return format!("DNS:{name}")
    } else if let Some(name) = name.uri() {
        return format!("URI:{name}")
    } else if let Some(ip) = name.ipaddress() {
        match ip.len() {
            4 => format!("IP:{}", Ipv4Addr::from(<[u8;4]>::try_from(ip).unwrap())),
            16 => format!("IP:{}", Ipv6Addr::from(<[u8;16]>::try_from(ip).unwrap())),
            _ => format!("IPx:{}", hex::encode(ip))
        }
    } else if let Some(dn) = name.directory_name() {
        format!("DN:{}", format_x509_name(dn))
    } else {
        "UNKNOWN".to_string()
    }
}
impl Reporter {
    pub async fn report_issuers(&self, issuers: &StackRef<X509>) {
        if !self.collect_digests {
            return;
        }
        // TODO This is always in whatever order the client returns the certificates :-/
        for issuer in issuers.iter() {
            self.note_issuer(issuer).await;
        }
    }

    async fn note_issuer(&self, x509: &X509Ref) -> Option<()> {
        let der_digest = x509.digest(MessageDigest::sha256()).ok()?.to_vec();
        // We try with a read lock first in order to increase parallelism.
        // Most of the time, the cert will already be in the store.
        let already = self.digests.read().await.contains(der_digest.as_slice());
        if !already {
            if self.digests.write().await.insert(der_digest) {
                if self.issuer_chan.send(x509.to_owned()).await.is_err() {

                }
            }
        }
        Some(())
    }

    pub async fn report_probe(&self, report: ProbeReport) -> Result<(), ReportingError> {
        info!(ip=%report.host, "Received report");
        if self.report_chan.send(report).await.is_err() {
            error!("Report formatter has exited early");
            Err(ReportingError::ReportFormatterFailed)
        } else {
            Ok(())
        }
    }
}

fn start_json(config: &JsonConfig) -> anyhow::Result<(impl Future<Output=()>+Send, Reporter)> {
    let (issuer_send, mut issuer_recv) = mpsc::channel::<X509>(5);
    let (report_send, mut report_recv) = mpsc::channel(5);

    let report_file = tokio::fs::File::from_std(std::fs::File::create(&config.output_file)?);
    let issuer_writer = config.issuer_file.as_ref().map(std::fs::File::create).transpose()?.map(tokio::fs::File::from_std);
    let has_issuer = issuer_writer.is_some();
    let container = config.container;
    let issuer_fut = async move {
        if let Some(issuer_file) = issuer_writer {
            let mut first_record = true;
            let mut issuer_file = tokio::io::BufWriter::new(issuer_file);
            if container {
                issuer_file.write_u8(b'[').await?;
            }
            while let Some(issuer) = issuer_recv.recv().await {
                match CertInfo::extract(issuer.as_ref()) {
                    Ok(info) => {
                        if container && !first_record {
                            issuer_file.write_u8(b',').await?;
                            first_record = false;
                        }
                        let info = serde_json::to_vec(&info)?;
                        issuer_file.write_all(info.as_slice()).await?;
                        issuer_file.write_u8(b'\n').await?;
                    }
                    Err(err) => {
                        warn!(%err, "Failed to extract data from certificate")
                    }
                }
            }
            if container {
                issuer_file.write_all(b"]").await?;
            }
            issuer_file.flush().await?;
        }
        Ok(()) as anyhow::Result<()>
    };

    let report_fut = async move {
        let mut file = BufWriter::new(report_file);
        if container {
            file.write_u8(b'[').await?;
        }
        let mut first_record = true;
        while let Some(report) = report_recv.recv().await {
            if container && !first_record {
                file.write_u8(b',').await?;
                first_record = false;
            }
            let json = serde_json::to_vec(&report)?;
            file.write_all(json.as_slice()).await?;
            file.write_u8(b'\n').await?;
        }
        if container {
            file.write_u8(b']').await?;
        }
        file.flush().await?;
        Ok(()) as anyhow::Result<()>
    };

    let task = async move {
        let (reporter, issuer) = futures::future::join(report_fut, issuer_fut).await;
        if let Err(error) = reporter {
            error!(%error, "Report writer failed")
        }
        if let Err(error) = issuer {
            error!(%error, "Issuer certificate writer failed")
        }
    };

    let reporter = Reporter {
        issuer_chan: issuer_send,
        report_chan: report_send,
        digests: Arc::new(Default::default()),
        collect_digests: has_issuer,
    };
    Ok((task, reporter))
}


/// Configure the reporting backend
pub(crate) fn configure_backend(config: &OutputFormat) -> anyhow::Result<(impl Future<Output=()>+Send, Reporter)> {
    match config {
        OutputFormat::Json(json) => start_json(json)
    }
}