peazyrsa/src/main.rs

use anyhow::{anyhow, Context, Result};
use clap::Parser;
use encoding::{label::encoding_from_whatwg_label, EncoderTrap};
use regex::Regex;
use std::{
    collections::BTreeMap,
    path::{Path, PathBuf},
    pin::Pin,
    sync::Arc,
};
use tokio::{
    fs::{self, File},
    io::{AsyncBufReadExt, BufReader},
};
use tokio::{pin, process::Command};

use async_stream::stream;

use futures_core::stream::Stream;
use futures_util::stream::StreamExt;

#[derive(Parser)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// new client name
    name: String,

    /// pki directory
    #[arg(short, long)]
    directory: Option<String>,

    /// client email
    #[arg(short, long)]
    email: Option<String>,

    /// files encoding
    #[arg(short = 'c', long)]
    encoding: Option<String>,

    /// keys subdir
    #[arg(long, default_value = "keys")]
    keys_dir: String,

    /// config subdir
    #[arg(long, default_value = "config")]
    config_dir: String,

    /// valid days
    #[arg(long, default_value = "30650")]
    days: u32,

    /// openssl binary
    #[arg(long, default_value = "openssl")]
    openssl: String,

    /// template file
    #[arg(long, default_value = "template.ovpn")]
    template_file: String,
}

struct VarsFile {
    filepath: PathBuf,
    vars: Option<BTreeMap<String, String>>,
    encoding: String,
}

struct AppConfig {
    encoding: String,
    req_days: u32,
    keys_subdir: String,
    config_subdir: String,
    template_file: String,
    openssl_default_cnf: String,
    openssl_cnf_env: String,
    ca_filename: String,
    default_email_domain: String,
    openssl: String,
    base_directory: String,
    email: String,
    name: String,
}

impl Default for AppConfig {
    fn default() -> Self {
        Self {
            encoding: "cp866".into(),
            req_days: 30650,
            keys_subdir: "keys".into(),
            config_subdir: "config".into(),
            template_file: "template.ovpn".into(),
            openssl_default_cnf: "openssl-1.0.0.cnf".into(),
            openssl_cnf_env: "KEY_CONFIG".into(),
            ca_filename: "ca.crt".into(),
            default_email_domain: "example.com".into(),
            openssl: "openssl".into(),
            base_directory: ".".into(),
            email: "name@example.com".into(),
            name: "user".into(),
        }
    }
}

impl From<&Args> for AppConfig {
    fn from(args: &Args) -> Self {
        let defaults = Self::default();

        let base_directory = args
            .directory
            .as_ref()
            .unwrap_or(&defaults.base_directory)
            .clone();
        let email = args.email.clone().unwrap_or(format!(
            "{}@{}",
            &args.name,
            defaults.default_email_domain.clone()
        ));
        let encoding = if let Some(enc) = args.encoding.clone() {
            enc.to_string()
        } else {
            defaults.encoding.clone()
        };
        let name = args.name.clone();
        let openssl = args.openssl.clone();
        let template_file = args.template_file.clone();
        let req_days = args.days;
        let keys_subdir = args.keys_dir.clone();
        let config_subdir = args.config_dir.clone();

        Self {
            base_directory,
            email,
            encoding,
            name,
            openssl,
            template_file,
            req_days,
            keys_subdir,
            config_subdir,
            ..defaults
        }
    }
}

async fn is_file_exist(filepath: &PathBuf) -> bool {
    let metadata = tokio::fs::metadata(&filepath).await;
    if metadata.is_err() {
        return false;
    }

    if !metadata.unwrap().is_file() {
        return false;
    }

    true
}

async fn read_file<'a, S, P>(filepath: P, encoding: S) -> Result<String>
where
    S: AsRef<str> + std::cmp::PartialEq<&'a str>,
    P: AsRef<Path>,
{
    let filepath = PathBuf::from(filepath.as_ref());
    if encoding == "utf8" {
        return Ok(fs::read_to_string(filepath).await?);
    }

    let enc = encoding_from_whatwg_label(encoding.as_ref()).ok_or(anyhow!("encoding not found"))?;

    let bytes = fs::read(filepath).await?;
    enc.decode(&bytes, encoding::DecoderTrap::Ignore)
        .map_err(|_| anyhow!("could not read file"))
}

async fn write_file(filepath: &PathBuf, text: String, encoding: &str) -> Result<()> {
    if encoding == "utf8" {
        return Ok(fs::write(filepath, text).await?);
    }

    let enc = encoding_from_whatwg_label(encoding).ok_or(anyhow!("encoding not found"))?;
    let mut bytes = Vec::new();
    enc.encode_to(&text, EncoderTrap::Ignore, &mut bytes)
        .map_err(|_| anyhow!("can't encode"))?;

    fs::write(filepath, bytes).await.context("can't write file")
}

async fn read_file_by_lines(
    filepath: &PathBuf,
    encoding: &str,
) -> Result<Box<dyn Stream<Item = String>>> {
    Ok(if encoding == "utf8" {
        let f = File::open(filepath).await?;
        let reader = BufReader::new(f);
        let mut lines = reader.lines();
        Box::new(stream! {
            while let Ok(Some(line)) = lines.next_line().await {
                yield line
            }
        })
    } else {
        let text = read_file(filepath, encoding).await?;
        Box::new(stream! {
            for line in text.lines() {
                yield line.to_string()
            }
        })
    })
}

impl VarsFile {
    async fn from_file(filepath: &PathBuf, encoding: String) -> Result<Self> {
        let metadata = tokio::fs::metadata(&filepath).await.context(format!(
            "file not found {}",
            filepath.to_str().expect("str")
        ))?;
        if !metadata.is_file() {
            Err(anyhow!("{} is not a file", filepath.to_str().expect("str")))?
        }
        Ok(VarsFile {
            filepath: filepath.to_path_buf(),
            vars: None,
            encoding,
        })
    }

    async fn from_dir(dir: PathBuf, encoding: String) -> Result<Self> {
        let filepath = dir.join("vars");
        let err_context = format!(
            "vars or vars.bat file not found in {}",
            dir.to_str().expect("str")
        );

        match Self::from_file(&filepath, encoding.clone()).await {
            Ok(res) => Ok(res),
            Err(_) => Self::from_file(&filepath.with_extension("bat"), encoding.clone())
                .await
                .map_err(|e| e.context(err_context)),
        }
    }

    async fn from_config(config: &AppConfig) -> Result<Self> {
        Self::from_dir(
            PathBuf::from(&config.base_directory),
            config.encoding.clone(),
        )
        .await
    }

    async fn parse(&mut self) -> Result<()> {
        let mut result = BTreeMap::new();
        let lines = read_file_by_lines(&self.filepath, &self.encoding).await?;
        let lines = Pin::from(lines);
        pin!(lines);

        let re_v2 =
            Regex::new(r#"^(export|set)\s\b(?P<key>[\w\d_]+)\b=\s?"?(?P<value>[^\#]+?)"?$"#)
                .context("regex v2")?;
        let re_v3 = Regex::new(r"^set_var\s(?P<key1>[\w\d_]+)\s+(?P<value1>[^\#]+?)$")
            .context("regex v3")?;

        while let Some(line) = lines.next().await {
            if let Some(caps) = re_v2.captures(line.as_str()) {
                result.insert(caps["key"].to_string(), caps["value"].to_string());
                continue;
            }

            if let Some(caps) = re_v3.captures(line.as_str()) {
                result.insert(caps["key"].to_string(), caps["value"].to_string());
            };
        }

        self.vars = Some(result);
        Ok(())
    }

    #[allow(dead_code)]
    fn apply(&self) -> Result<()> {
        if let Some(vars) = self.vars.clone() {
            for (key, value) in vars.iter() {
                unsafe {
                    std::env::set_var(key, value);
                }
            }
        } else {
            Err(anyhow!("vars not parsed"))?
        }

        Ok(())
    }
}

trait ICryptoProvider {
    async fn request(&self) -> Result<()>;
    async fn sign(&self) -> Result<()>;
}

struct OpenSSLProvider {
    vars: BTreeMap<String, String>,
    base_dir: PathBuf,
    openssl_cnf: PathBuf,
    openssl: String,
    ca_file: PathBuf,
    req_file: PathBuf,
    key_file: PathBuf,
    cert_file: PathBuf,
    req_days: u32,
}

impl OpenSSLProvider {
    async fn is_ca_exists(&self) -> bool {
        is_file_exist(&self.ca_file).await
    }

    async fn is_cert_exists(&self) -> bool {
        is_file_exist(&self.cert_file).await
    }

    async fn is_req_exists(&self) -> bool {
        is_file_exist(&self.req_file).await
    }

    fn from_cfg(cfg: &AppConfig, vars: BTreeMap<String, String>) -> Self {
        let base_dir = PathBuf::from(&cfg.base_directory);
        let keys_dir = base_dir.clone().join(cfg.keys_subdir.clone());
        let name = cfg.name.clone();
        let mut vars = vars;

        vars.insert("KEY_CN".into(), name.clone());
        vars.insert("KEY_NAME".into(), name.clone());
        vars.insert("KEY_EMAIL".into(), cfg.email.clone());

        let ca_file = keys_dir.join(cfg.ca_filename.clone());
        let req_file = keys_dir.join(format!("{}.csr", &name));
        let key_file = keys_dir.join(format!("{}.key", &name));
        let cert_file = keys_dir.join(format!("{}.crt", &name));
        let openssl_cnf = base_dir.clone().join(
            std::env::var(cfg.openssl_cnf_env.clone()).unwrap_or(cfg.openssl_default_cnf.clone()),
        );

        Self {
            vars,
            base_dir,
            openssl_cnf,
            openssl: cfg.openssl.clone(),
            ca_file,
            req_file,
            key_file,
            cert_file,
            req_days: cfg.req_days,
        }
    }
}

impl ICryptoProvider for OpenSSLProvider {
    async fn request(&self) -> Result<()> {
        if self.is_req_exists().await {
            return Ok(());
        }

        if !self.is_ca_exists().await {
            return Err(anyhow!(
                "ca file not found: {}",
                &self.ca_file.to_str().unwrap()
            ));
        }

        let status = Command::new(&self.openssl)
            .args([
                "req",
                "-nodes",
                "-new",
                "-keyout",
                self.key_file.to_str().unwrap(),
                "-out",
                self.req_file.to_str().unwrap(),
                "-config",
                self.openssl_cnf.to_str().unwrap(),
                "-batch",
            ])
            .current_dir(&self.base_dir)
            .envs(&self.vars)
            .status()
            .await?;

        match status.success() {
            true => Ok(()),
            false => Err(anyhow!("openssl req execution failed")),
        }
    }

    async fn sign(&self) -> Result<()> {
        if self.is_cert_exists().await {
            return Ok(());
        }

        if !self.is_ca_exists().await {
            return Err(anyhow!(
                "ca file not found: {}",
                &self.ca_file.to_str().unwrap()
            ));
        }

        let status = Command::new(&self.openssl)
            .args([
                "ca",
                "-days",
                format!("{}", self.req_days).as_str(),
                "-out",
                self.cert_file.to_str().unwrap(),
                "-in",
                self.req_file.to_str().unwrap(),
                "-config",
                self.openssl_cnf.to_str().unwrap(),
                "-batch",
            ])
            .current_dir(&self.base_dir)
            .envs(&self.vars)
            .status()
            .await?;

        match status.success() {
            true => Ok(()),
            false => Err(anyhow!("ssl ca execution failed")),
        }
    }
}

struct Certs<T>
where
    T: ICryptoProvider,
{
    encoding: String,
    ca_file: PathBuf,
    key_file: PathBuf,
    cert_file: PathBuf,
    config_file: PathBuf,
    template_file: PathBuf,
    provider: Arc<T>,
}

fn make_certs_provider(cfg: &AppConfig, vars: BTreeMap<String, String>) -> impl ICryptoProvider {
    OpenSSLProvider::from_cfg(cfg, vars)
}

impl<T> Certs<T>
where
    T: ICryptoProvider,
{
    fn new(cfg: &AppConfig, provider: T) -> Self {
        let base_dir = PathBuf::from(&cfg.base_directory);
        let keys_dir = base_dir.clone().join(cfg.keys_subdir.clone());
        let config_dir = base_dir.clone().join(cfg.config_subdir.clone());
        let name = cfg.name.clone();

        Certs {
            encoding: cfg.encoding.clone(),
            ca_file: keys_dir.join(cfg.ca_filename.clone()),
            key_file: keys_dir.join(format!("{}.key", &name)),
            cert_file: keys_dir.join(format!("{}.crt", &name)),
            config_file: config_dir.join(format!("{}.ovpn", &name)),
            template_file: base_dir.clone().join(cfg.template_file.clone()),
            provider: Arc::new(provider),
        }
    }

    async fn is_config_exists(&self) -> bool {
        is_file_exist(&self.config_file).await
    }

    async fn request(&self) -> Result<()> {
        self.provider.request().await
    }

    async fn sign(&self) -> Result<()> {
        self.provider.sign().await
    }

    async fn build_client_config(&self) -> Result<bool> {
        if self.is_config_exists().await {
            return Ok(false);
        }

        self.request().await?;
        self.sign().await?;

        let (template_file, ca_file, cert_file, key_file) = (
            self.template_file.clone(),
            self.ca_file.clone(),
            self.cert_file.clone(),
            self.key_file.clone(),
        );
        let enc = self.encoding.clone();
        let (enc1, enc2, enc3, enc4) = (enc.clone(), enc.clone(), enc.clone(), enc.clone());

        if let (Ok(Ok(template)), Ok(Ok(ca)), Ok(Ok(cert)), Ok(Ok(key))) = tokio::join!(
            tokio::spawn(read_file(template_file, enc1)),
            tokio::spawn(read_file(ca_file, enc2)),
            tokio::spawn(read_file(cert_file, enc3)),
            tokio::spawn(read_file(key_file, enc4))
        ) {
            let text = template
                .replace("{{ca}}", ca.trim())
                .replace("{{cert}}", cert.trim())
                .replace("{{key}}", key.trim());

            write_file(&self.config_file, text, &self.encoding).await?;

            Ok(true)
        } else {
            Err(anyhow!("files read error"))
        }
    }
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
    let args = Args::parse();
    let config = AppConfig::from(&args);
    let mut vars = VarsFile::from_config(&config).await?;
    vars.parse().await?;

    println!("found vars: {}", vars.filepath.to_str().expect("fff"));
    println!("loaded: {:#?}", &vars.vars);

    let provider = make_certs_provider(&config, vars.vars.unwrap());
    let certs = Certs::new(&config, provider);
    let created = certs.build_client_config().await?;

    let result_file = certs.config_file.to_str().unwrap();

    if created {
        println!("created: {result_file}");
        Ok(())
    } else {
        Err(anyhow!("file exists: {result_file}"))
    }
}