From 540a11f8c4d3f3f449c935f7030823522a120651 Mon Sep 17 00:00:00 2001
From: Blallo <blallo@autistici.org>
Date: Wed, 19 Aug 2020 18:43:51 +0200
Subject: [PATCH] Updated gen_keys script

Improved cli. Now it also allows to use tagged vaults.
---
 keys/gen_keys.py | 105 ++++++++++++++++++++++++++++++++++-------------
 1 file changed, 76 insertions(+), 29 deletions(-)

diff --git a/keys/gen_keys.py b/keys/gen_keys.py
index 456ca9c..644a736 100755
--- a/keys/gen_keys.py
+++ b/keys/gen_keys.py
@@ -1,5 +1,7 @@
 #!/usr/bin/env python
 
+import argparse
+import os
 import subprocess
 import sys
 import typing as T
@@ -47,45 +49,90 @@ def gen_key(name: T.Text) -> None:
         subprocess.call(["wg", "pubkey"], stdin=sec_r, stdout=pub)
 
 
-def to_vault(name: T.Text, passfile: T.Text) -> None:
+def _encrypt_string(
+    plaintext: T.Text,
+    passfile: T.Text,
+    vault_id: T.Optional[T.Text],
+    passfile_dir: T.Optional[T.Text],
+) -> bytes:
+    if passfile_dir is None:
+        passfile_dir = os.getcwd()
+    if vault_id is None:
+        return _encrypt_string_simple(plaintext, passfile, passfile_dir)
+    else:
+        return _encrypt_string_vault(plaintext, f"{vault_id}@{passfile}", passfile_dir)
+
+
+def _encrypt_string_simple(
+    plaintext: T.Text, passfile: T.Text, passfile_dir: T.Text
+) -> bytes:
+    return subprocess.check_output(
+        [
+            "ansible-vault",
+            "encrypt_string",
+            f"--vault-password-file={passfile}",
+            plaintext,
+        ],
+        cwd=passfile_dir,
+    )
+
+
+def _encrypt_string_vault(
+    plaintext: T.Text, vault_passfile: T.Text, passfile_dir: T.Text
+) -> bytes:
+    return subprocess.check_output(
+        ["ansible-vault", "encrypt_string", f"--vault-id={vault_passfile}", plaintext,],
+        cwd=passfile_dir,
+    )
+
+
+def to_vault(
+    name: T.Text,
+    passfile: T.Text,
+    vault_id: T.Optional[T.Text],
+    passfile_dir: T.Optional[T.Text],
+) -> None:
     with open(f"{name}.pub", "r") as pub:
-        pubkey = pub.readline()
-        enc_pub = subprocess.check_output(
-            [
-                "ansible-vault",
-                "encrypt_string",
-                f"--vault-password-file={passfile}",
-                pubkey,
-            ]
-        )
+        pubkey = pub.readline().strip("\n")
+        enc_pub = _encrypt_string(pubkey, passfile, vault_id, passfile_dir)
     with open(f"{name}.sec", "r") as sec:
-        seckey = sec.readline()
-        enc_sec = subprocess.check_output(
-            [
-                "ansible-vault",
-                "encrypt_string",
-                f"--vault-password-file={passfile}",
-                seckey,
-            ]
-        )
+        seckey = sec.readline().strip("\n")
+        enc_sec = _encrypt_string(seckey, passfile, vault_id, passfile_dir)
     HOSTS[name] = {
         "public_key": yaml.load(enc_pub.decode(ENCODING), Loader=yaml.SafeLoader),
         "private_key": yaml.load(enc_sec.decode(ENCODING), Loader=yaml.SafeLoader),
     }
 
 
-def usage() -> None:
-    print("Usage: \n\tgen_keys.py <path/to/inventory> <path/to/vault_password_file>")
-
-
 if __name__ == "__main__":
-    if len(sys.argv) != 3:
-        usage()
-        sys.exit(-1)
-
-    for host in load_hosts(sys.argv[1]):
+    parser = argparse.ArgumentParser(
+        description="Generates keys for wireguard connection of the dns servers"
+    )
+    parser.add_argument(
+        "inventory", metavar="INVENTORY", type=str, help="path to the inventory"
+    )
+    parser.add_argument(
+        "passfile",
+        metavar="VAULT_PASSFILE",
+        type=str,
+        help="the name of the file that contains the passphrase for the inventory",
+    )
+    parser.add_argument(
+        "--vault-id",
+        metavar="VAULT_ID",
+        type=str,
+        help="the name of the (existing) vault",
+    )
+    parser.add_argument(
+        "--passfile-dir",
+        metavar="PASSFILE_DIR",
+        type=str,
+        help="path where the passfile is located",
+    )
+    args = parser.parse_args()
+    for host in load_hosts(args.inventory):
         gen_key(host)
-        to_vault(host, sys.argv[2])
+        to_vault(host, args.passfile, args.vault_id, args.passfile_dir)
         result = yaml.dump(HOSTS, Dumper=yaml.SafeDumper)
         with open("result.yml", "w") as res:
             res.write(result)