// SPDX-License-Identifier: Apache-2.0

use crate::config::Config;
use crate::snp::grpc::proto;
use log::debug;
use std::net::IpAddr;

#[derive(Debug)]
pub struct Args {
    pub uuid: String,
    pub vcpus: u32,
    pub kernel: String,
    pub initrd: String,
    pub args: proto::MeasurementArgs,
}

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("Failed to execute: {0}")]
    FailedExecution(String),
    #[error("Disk access error: {0}")]
    DiskAccess(#[from] std::io::Error),
    #[error("Could not parse kernel parameters sent by deamon")]
    AlteredKernelParams,
    #[error("Could not inject secrets to VM: {0}")]
    InjectorError(String),
    #[error(transparent)]
    ConfigError(#[from] crate::config::Error),
    #[error("The endpoint returned by the daemon is not valid: {0}")]
    InvalidVmAddress(String),
}

impl From<std::net::AddrParseError> for Error {
    fn from(_: std::net::AddrParseError) -> Self {
        Self::AlteredKernelParams
    }
}

impl From<std::num::ParseIntError> for Error {
    fn from(_: std::num::ParseIntError) -> Self {
        Self::AlteredKernelParams
    }
}

impl Args {
    pub fn get_measurement(&self) -> Result<String, Error> {
        log::debug!("Measurement args: {self:#?}");
        let params = self.kernel_params()?;
        Config::verify_and_install_artefacts(&self.args.ovmf_hash)?;
        Config::verify_and_install_artefacts(&self.initrd)?;
        Config::verify_and_install_artefacts(&self.kernel)?;
        let artefacts_dir = Config::artefacts_dir()? + "/";
        let script_result = std::process::Command::new("sev-snp-measure.py")
            .arg("--mode=snp")
            .arg(format!("--vcpus={}", self.vcpus))
            .arg("--vcpu-type=EPYC-v4")
            .arg("--output-format=hex")
            .arg(format!("--ovmf={}", artefacts_dir.clone() + &self.args.ovmf_hash))
            .arg(format!("--kernel={}", artefacts_dir.clone() + &self.kernel))
            .arg(format!("--initrd={}", artefacts_dir + &self.initrd))
            .arg(format!("--append={}", params))
            .output()
            .map_err(|e| {
                Error::FailedExecution(format!("Failed to execute sev-snp-measure.py: {e:?}"))
            })?;
        if !script_result.status.success() {
            return Err(Error::FailedExecution(format!(
                "sev-snp-measure.py failed: !!! stdout:\n{}\n!!! stderr:\n{}",
                String::from_utf8(script_result.stdout.clone())
                    .unwrap_or("Could not grab stdout from installation script.".to_string()),
                String::from_utf8(script_result.stderr.clone())
                    .unwrap_or("Could not grab stderr from installation script.".to_string())
            )));
        }
        Ok(String::from_utf8(script_result.stdout)
            .map_err(|e| {
                Error::FailedExecution(format!("Could not parse sev-snp-measure.py output: {e:?}"))
            })?
            .trim_end()
            .to_string())
    }

    // This must produce the exact same params as the same operation on the daemon.
    // Passing kernel params as String from the daemon leaves the install open for attacks.
    fn kernel_params(&self) -> Result<String, Error> {
        let mut ip_string = String::new();
        let mut public_ipv4 = false;
        for ip in self.args.ips.iter() {
            if ip.address.parse::<IpAddr>()?.is_ipv4() {
                public_ipv4 = true;
            }
            ip.gateway.parse::<IpAddr>()?;
            ip.mask.parse::<u128>()?;
            ip_string += &format!(
                "detee_net_eth{}={}_{}_{} ",
                ip.nic_index, ip.address, ip.mask, ip.gateway
            );
        }
        if !public_ipv4 {
            ip_string = "detee_net_eth0=10.0.2.15_24_10.0.2.2 ".to_string() + &ip_string;
        }
        let admin_key = format!("detee_admin={} ", Config::get_detee_wallet()?);
        let hostname = format!("detee_uuid={}", self.uuid);
        let params = format!("{}{}{}", ip_string, admin_key, hostname);
        debug!("Calculated kernel params for {} to: {}", self.uuid, params);
        Ok(params)
    }
}

pub fn execute(
    measurement: String,
    server_addr: String,
    os_template: Option<(&str, &str)>,
    vm_hostname: &str,
) -> Result<super::VmSshArgs, Error> {
    let parsed_addr = match server_addr.parse::<std::net::SocketAddrV4>() {
        Ok(addr) => addr,
        Err(_) => return Err(Error::InvalidVmAddress(server_addr)),
    };
    let ssh_pubkey = Config::init_config().get_ssh_pubkey()?;
    let ssh_args = super::VmSshArgs {
        ip: parsed_addr.ip().to_string(),
        port: parsed_addr.port().to_string(),
        user: "root".to_string(),
        key_path: ssh_pubkey.trim_end_matches(".pub").to_string(),
        ..Default::default()
    };
    eprintln!("Injecting disk encryption key into VM. This will take a minute. Do not interrupt.");
    let (os_template_url, os_template_sha) = os_template.unwrap_or(("", ""));
    let logs_path = Config::logs_dir()? + "/" + vm_hostname;
    log::info!("Logs will be saved to {}", logs_path);
    let logs_file = std::fs::File::create(logs_path.clone())?;
    let mut child_process = std::process::Command::new("detee-cli_injector.sh")
        .env("SERVER_ADDR", server_addr)
        .env("SSH_KEY_FILE", ssh_pubkey)
        .env("DETEE_INSTALL_URL", os_template_url)
        .env("DETEE_INSTALL_SHA", os_template_sha)
        .env("MEASUREMENT", measurement)
        .env("VM_HOSTNAME", vm_hostname)
        .stdout(logs_file.try_clone()?)
        .stderr(logs_file)
        .spawn()?;
    let status = child_process.wait()?;
    if !status.success() {
        log::error!(
            "detee-cli_injector.sh exit code: {:?}! Check logs at: {}",
            status.code(),
            logs_path
        );
        return Err(Error::InjectorError("Script execution failed".to_string()));
    }
    Ok(ssh_args)
}