From b701deb376619e3abd4dc0e62268f987b5639411 Mon Sep 17 00:00:00 2001
From: Zhiguang Wang <wangzhiguang@kingsoft.com>
Date: Sun, 29 Apr 2018 12:26:46 +0800
Subject: [PATCH] Refactoring.

---
 main.py | 51 +++++++++++++++++++++++++--------------------------
 1 file changed, 25 insertions(+), 26 deletions(-)

diff --git a/main.py b/main.py
index 8f3592f..808d1ab 100644
--- a/main.py
+++ b/main.py
@@ -11,7 +11,14 @@
 # Temp dir of Lambda runtime
 CERTBOT_DIR = '/tmp/certbot'
 
-def call_certbot(email, domains):
+def rm_tmp_dir():
+    if os.path.exists(CERTBOT_DIR):
+        try:
+            shutil.rmtree(CERTBOT_DIR)
+        except NotADirectoryError:
+            os.remove(CERTBOT_DIR)
+
+def obtain_certs(email, domains):
     certbot_args = [
         # Override directory paths so script doesn't have to be run as root
         '--config-dir', CERTBOT_DIR,
@@ -29,25 +36,24 @@ def call_certbot(email, domains):
     ]
     return certbot.main.main(certbot_args)
 
-def publish_certs_to_s3(bucket, prefix):
+# /tmp/certbot
+# ├── live
+# │   └── [domain]
+# │       ├── README
+# │       ├── cert.pem
+# │       ├── chain.pem
+# │       ├── fullchain.pem
+# │       └── privkey.pem
+def upload_certs(s3_bucket, s3_prefix):
     client = boto3.client('s3')
-
-    # /tmp/certbot
-    # ├── live
-    # │   └── [domain]
-    # │       ├── README
-    # │       ├── cert.pem
-    # │       ├── chain.pem
-    # │       ├── fullchain.pem
-    # │       └── privkey.pem
     cert_dir = os.path.join(CERTBOT_DIR, 'live')
     for dirpath, _dirnames, filenames in os.walk(cert_dir):
         for filename in filenames:
             local_path = os.path.join(dirpath, filename)
             relative_path = os.path.relpath(local_path, cert_dir)
-            key = os.path.join(prefix, relative_path)
-            print(f'Uploading: {local_path} => s3://{bucket}/{key}')
-            client.upload_file(local_path, bucket, key)
+            s3_key = os.path.join(s3_prefix, relative_path)
+            print(f'Uploading: {local_path} => s3://{s3_bucket}/{s3_key}')
+            client.upload_file(local_path, s3_bucket, s3_key)
 
 def guarded_handler(event, context):
     # Input parameters
@@ -56,21 +62,14 @@ def guarded_handler(event, context):
     s3_bucket = event['s3_bucket']  # The S3 bucket to publish certificates
     s3_prefix = event['s3_prefix']  # The S3 key prefix to publish certificates
 
-    call_certbot(email, domains)
-    publish_certs_to_s3(s3_bucket, s3_prefix)
+    obtain_certs(email, domains)
+    upload_certs(s3_bucket, s3_prefix)
 
-    return f'Certs published successfully.'
-
-def cleanup():
-    if os.path.exists(CERTBOT_DIR):
-        try:
-            shutil.rmtree(CERTBOT_DIR)
-        except NotADirectoryError:
-            os.remove(CERTBOT_DIR)
+    return 'Certificates obtained and uploaded successfully.'
 
 def lambda_handler(event, context):
     try:
-        cleanup()
+        rm_tmp_dir()
         return guarded_handler(event, context)
     finally:
-        cleanup()
+        rm_tmp_dir()