acme-client-rs/src/main.rs

#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate rust_util;

use rust_util::XResult;
use acme_lib::{DirectoryUrl, Directory, create_p384_key, create_p256_key, create_rsa_key};
use acme_lib::persist::FilePersist;
use clap::{App, Arg};
use std::sync::RwLock;
use std::collections::BTreeMap;
use tide::Request;
use std::process::exit;
use std::time::Duration;
use async_std::channel::Sender;

const NAME: &str = env!("CARGO_PKG_NAME");
const VERSION: &str = env!("CARGO_PKG_VERSION");
const AUTHORS: &str = env!("CARGO_PKG_AUTHORS");
const DESCRIPTION: &str = env!("CARGO_PKG_DESCRIPTION");

lazy_static! {
    static ref TOKEN_MAP: RwLock<BTreeMap<String, String>> = RwLock::new(BTreeMap::new());
}

#[derive(Debug, Clone, Copy)]
enum Algo {
    Ec256,
    Ec384,
    Rsa(u32),
}

#[derive(Debug, Clone, Copy)]
enum Mode {
    Prod,
    Test,
}

#[async_std::main]
async fn main() -> tide::Result<()> {
    println!("{}", include_str!("logo.txt"));

    let matches = App::new(NAME)
        .version(VERSION)
        .about(DESCRIPTION)
        .author(AUTHORS)
        .arg(Arg::with_name("version").short("V").long("version").help("Print version"))
        .arg(Arg::with_name("verbose").short("v").long("verbose").help("Verbose"))
        .arg(Arg::with_name("type").short("t").long("type").default_value("http").takes_value(true).help("Type http or dns"))
        .arg(Arg::with_name("port").short("p").long("port").default_value("80").takes_value(true).help("Http port"))
        .arg(Arg::with_name("domain").short("d").long("domain").multiple(true).takes_value(true).help("Domains"))
        .arg(Arg::with_name("email").long("email").takes_value(true).help("Contract email"))
        .arg(Arg::with_name("algo").short("a").long("algo").takes_value(true).default_value("ec384").help("Pki algo"))
        .arg(Arg::with_name("timeout").long("timeout").takes_value(true).default_value("5000").help("Timeout (ms)"))
        .arg(Arg::with_name("mode").short("m").long("mode").takes_value(true).default_value("prod").help("Mode"))
        .arg(Arg::with_name("dir").long("dir").takes_value(true).default_value("acme_dir").help("Account key dir"))
        .get_matches();

    if matches.is_present("version") {
        information!("{} v{}", NAME, VERSION);
        exit(1);
    }

    let email = matches.value_of("email").unwrap_or_else(|| {
        failure!("Email is not assigned.");
        exit(1);
    });

    if let None = matches.value_of("type") {
        failure!("Type is not assigned.");
        exit(1);
    }
    let port: u16 = match matches.value_of("port") {
        Some(p) => p.parse().unwrap_or_else(|e| {
            failure!("Parse port: {}, failed: {}", p, e);
            exit(1);
        }),
        None => {
            failure!("Port is not assigned.");
            exit(1);
        }
    };
    let timeout: u64 = match matches.value_of("timeout") {
        Some(p) => p.parse().unwrap_or_else(|e| {
            failure!("Parse timeout: {}, failed: {}", p, e);
            exit(1);
        }),
        None => {
            failure!("Timeout is not assigned.");
            exit(1);
        }
    };
    let algo = match matches.value_of("algo") {
        Some("ec256") => Algo::Ec256,
        Some("ec384") => Algo::Ec384,
        Some("rsa2048") => Algo::Rsa(2048),
        Some("rsa3072") => Algo::Rsa(3072),
        Some("rsa4096") => Algo::Rsa(4096),
        _ => {
            failure!("Algo is not assigned, or wrong, should be: ec256, ec384, rsa2048, rsa3073 or rsa4096.");
            exit(1);
        }
    };
    let mode = match matches.value_of("mode") {
        Some("prod") => Mode::Prod,
        Some("test") => Mode::Test,
        _ => {
            failure!("Mode is not assigned, or wrong, should be: prod or test");
            exit(1);
        }
    };
    let dir = matches.value_of("dir").unwrap_or("acme_dir");

    let domains_val = matches.values_of("domain").unwrap_or_else(|| {
        failure!("Domains is not assigned.");
        exit(1);
    });

    let (s, r) = async_std::channel::bounded(1);
    startup_http_server(s, port);
    r.recv().await.ok();
    async_std::task::sleep(Duration::from_millis(500)).await;

    let domains: Vec<&str> = domains_val.collect::<Vec<_>>();
    let primary_name = domains[0];
    let alt_names: Vec<&str> = domains.into_iter().skip(1).collect();
    information!("Domains, main: {}, alt: {:?}", primary_name, alt_names);
    let acme_request = AcmeRequest {
        contract_email: email,
        primary_name,
        alt_names: &alt_names,
        algo,
        mode,
        dir,
        timeout,
    };
    if let Err(e) = request_domains(acme_request) {
        failure!("Request certificate by acme failed: {}", e);
        exit(1);
    }
    Ok(())
}

#[derive(Debug)]
struct AcmeRequest<'a> {
    contract_email: &'a str,
    primary_name: &'a str,
    alt_names: &'a [&'a str],
    algo: Algo,
    mode: Mode,
    dir: &'a str,
    timeout: u64,
}

fn request_domains(acme_request: AcmeRequest) -> XResult<()> {
    information!("Acme mode: {:?}", acme_request.mode);
    let url = match acme_request.mode {
        Mode::Prod => DirectoryUrl::LetsEncrypt,
        Mode::Test => DirectoryUrl::LetsEncryptStaging,
    };
    information!("Acme dir: {}", acme_request.dir);
    std::fs::create_dir(acme_request.dir).ok();
    let persist = FilePersist::new(acme_request.dir);
    let dir = opt_result!(Directory::from_url(persist, url), "Create directory from url failed: {}");
    let acc = opt_result!(dir.account(acme_request.contract_email), "Directory set account failed: {}");
    let mut ord_new = opt_result!( acc.new_order(acme_request.primary_name, acme_request.alt_names), "Create order failed: {}");

    let mut order_csr_index = 0;
    let ord_csr = loop {
        if let Some(ord_csr) = ord_new.confirm_validations() {
            break ord_csr;
        }

        information!("Loop for acme challenge auth, #{}", order_csr_index);
        order_csr_index += 1;

        let auths = opt_result!(ord_new.authorizations(), "Order auth failed: {}");
        for auth in &auths {
            let chall = auth.http_challenge();
            let token = chall.http_token();
            // let path = format!(".well-known/acme-challenge/{}", token);
            let proof = chall.http_proof();

            {
                information!("Add acme challenge: {} -> {}",token, proof);
                TOKEN_MAP.write().unwrap().insert(token.to_string(), proof);
            }
            opt_result!(chall.validate(acme_request.timeout), "Validate challenge failed: {}");
        }

        opt_result!(ord_new.refresh(), "Refresh order failed: {}");
    };

    information!("Generate private key, type: {:?}", acme_request.algo);
    let pkey_pri = match acme_request.algo {
        Algo::Ec256 => create_p256_key(),
        Algo::Ec384 => create_p384_key(),
        Algo::Rsa(bits) => create_rsa_key(bits),
    };
    information!("Created private key: {:?}", pkey_pri);

    let ord_cert = opt_result!( ord_csr.finalize_pkey(pkey_pri, acme_request.timeout), "Submit CSR failed: {}");
    let cert = opt_result!( ord_cert.download_and_save_cert(), "Download and save certificate failed: {}");

    information!("Created certificate: {:?}", cert);

    Ok(())
}

fn startup_http_server(s: Sender<i32>, port: u16) {
    async_std::task::spawn(async move {
        information!("Listen at 0.0.0.0:{}", port);
        let mut app = tide::new();
        app.at("/.well-known/acme-challenge/:token").get(|req: Request<()>| async move {
            let token = match req.param("token") {
                Ok(token) => token,
                Err(e) => {
                    warning!("Cannot get token from url, query: {:?}, error: {}", req.url().query(), e);
                    return Ok("400 - bad request".to_string());
                }
            };
            let peer = req.peer_addr().unwrap_or("none");
            let auth_token = { TOKEN_MAP.read().unwrap().get(token).cloned() };
            match auth_token {
                Some(auth_token) => {
                    information!("Request acme challenge: {} -> {}, peer: {:?}", token, auth_token, peer);
                    Ok(auth_token)
                }
                None => {
                    warning!("Request acme challenge not found: {}, peer: {:?}", token, peer);
                    Ok("404 - not found".to_string())
                }
            }
        });
        s.send(1).await.ok();
        if let Err(e) = app.listen(&format!("0.0.0.0:{}", port)).await {
            failure!("Failed to listen 0.0.0.0:{}, program will exit, error: {}", port, e);
            exit(1);
        }
    });
}