From 52085597c1180225259071b3694def01e4e1767d Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sat, 1 May 2021 08:38:36 +0800
Subject: [PATCH] feat: use acme_request

---
 src/main.rs | 51 +++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 37 insertions(+), 14 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 2ae8282..7504c0b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -23,14 +23,14 @@ lazy_static! {
     static ref TOKEN_MAP: RwLock<BTreeMap<String, String>> = RwLock::new(BTreeMap::new());
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 enum Algo {
     Ec256,
     Ec384,
     Rsa(u32),
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 enum Mode {
     Prod,
     Test,
@@ -53,6 +53,7 @@ async fn main() -> tide::Result<()> {
         .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") {
@@ -108,6 +109,7 @@ async fn main() -> tide::Result<()> {
             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.");
@@ -123,24 +125,45 @@ async fn main() -> tide::Result<()> {
     let primary_name = domains[0];
     let alt_names: Vec<&str> = domains.into_iter().skip(1).collect();
     information!("Domains, main: {}, alt: {:?}", primary_name, alt_names);
-    if let Err(e) = request_domains(email, primary_name, &alt_names, algo, mode, timeout) {
+    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(())
 }
 
-fn request_domains(contract_email: &str, primary_name: &str, alt_names: &[&str], algo: Algo, mode: Mode, timeout: u64) -> XResult<()> {
-    information!("Acme mode: {:?}", mode);
-    let url = match mode {
+#[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,
     };
-    std::fs::create_dir("__temp_dir").ok();
-    let persist = FilePersist::new("__temp_dir");
+    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(contract_email), "Directory set account failed: {}");
-    let mut ord_new = opt_result!( acc.new_order(primary_name, alt_names), "Create order 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 {
@@ -162,21 +185,21 @@ fn request_domains(contract_email: &str, primary_name: &str, alt_names: &[&str],
                 information!("Add acme challenge: {} -> {}",token, proof);
                 TOKEN_MAP.write().unwrap().insert(token.to_string(), proof);
             }
-            opt_result!(chall.validate(timeout), "Validate challenge failed: {}");
+            opt_result!(chall.validate(acme_request.timeout), "Validate challenge failed: {}");
         }
 
         opt_result!(ord_new.refresh(), "Refresh order failed: {}");
     };
 
-    information!("Generate private key, type: {:?}", algo);
-    let pkey_pri = match algo {
+    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, timeout), "Submit CSR failed: {}");
+    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);