1 files changed, 111 insertions, 18 deletions

diff --git a/vmdb/plugins/cryptsetup_plugin.py b/vmdb/plugins/cryptsetup_plugin.py
index 238d7ab..acb9b04 100644
--- a/vmdb/plugins/cryptsetup_plugin.py
+++ b/vmdb/plugins/cryptsetup_plugin.py
@@ -30,34 +30,127 @@ class CryptsetupPlugin(vmdb.Plugin):
 
 class CryptsetupStepRunner(vmdb.StepRunnerInterface):
     def get_key_spec(self):
-        return {"cryptsetup": str, "password": str, "name": str}
+        return {
+            "cryptsetup": str,
+            "tag": str,
+            "password": "",
+            "key-file": "",
+            "key-cmd": "",
+        }
 
     def run(self, step, settings, state):
-        cleartext_tag = step["cryptsetup"]
-        password = step["password"]
-        name = step["name"]
+        underlying = values["cryptsetup"]
+        crypt_name = values["tag"]
+        password = values["password"] or None
+        key_file = values["key-file"] or None
+        key_cmd = values["key-cmd"] or None
 
-        device = state.tags.get_dev(cleartext_tag)
-        tmp = tempfile.mkdtemp()
-        key = os.path.join(tmp, "key")
-        with open(key, "w") as f:
-            f.write(password)
+        if not isinstance(underlying, str):
+            raise vmdb.NotString("cryptsetup", underlying)
+
+        if not isinstance(crypt_name, str):
+            raise vmdb.NotString("cryptsetup: tag", crypt_name)
+
+        if password is None and key_file is None and key_cmd is None:
+            raise Exception(
+                "cryptsetup step MUST define one of password, key-file, or key-cmd"
+            )
+
+        if password is not None and key_file is not None:
+            raise Exception(
+                "cryptsetup step MUST define only one of password or key-file"
+            )
+
+        if password is not None and key_cmd is not None:
+            raise Exception(
+                "cryptsetup step MUST define only one of password or key-cmd"
+            )
+
+        if key_file is not None and key_cmd is not None:
+            raise Exception(
+                "cryptsetup step MUST define only one of key_file or key-cmd"
+            )
+
+        state.tmp_key_file = None
+        rmtmp = False
+
+        if password is not None:
+            key_file = self._write_temp(password)
+            rmtp = True
+
+        if key_cmd is not None:
+            output = vmdb.runcmd(["sh", "-ec", key_cmd])
+            output = output.decode("UTF-8")
+            key = output.splitlines()[0]
+            key_file = self._write_temp(key)
+            rmtp = True
+
+        assert key_file is not None
+
+        dev = state.tags.get_dev(underlying)
+        if dev is None:
+            for t in state.tags.get_tags():
+                logging.debug(
+                    "tag %r dev %r mp %r",
+                    t,
+                    state.tags.get_dev(t),
+                    state.tags.get_builder_mount_point(t),
+                )
+            assert 0
 
-        vmdb.runcmd(["cryptsetup", "luksFormat", "--batch-mode", device, key])
         vmdb.runcmd(
-            ["cryptsetup", "open", "--type=luks", "--key-file", key, device, name]
+            [
+                "cryptsetup",
+                "luksFormat",
+                "--batch-mode",
+                "--type=luks2",
+                "--allow-discards",
+                dev,
+                key_file,
+            ]
         )
-        crypt_device = f"/dev/mapper/{name}"
-        assert os.path.exists(crypt_device)
+        vmdb.runcmd(
+            [
+                "cryptsetup",
+                "open",
+                "--key-file",
+                key_file,
+                dev,
+                crypt_name,
+            ]
+        )
+
+        crypt_dev = "/dev/mapper/{}".format(crypt_name)
+        assert os.path.exists(crypt_dev)
 
         uuid = vmdb.runcmd(["cryptsetup", "luksUUID", device]).decode("UTF8").strip()
 
-        state.tags.append(name)
-        state.tags.set_dev(name, crypt_device)
+        state.tags.append(crypt_name)
+        state.tags.set_dev(crypt_name, crypt_dev)
         state.tags.set_luksuuid(name, uuid)
         state.tags.set_dm(name, name)
-        vmdb.progress(f"LUKS: name={name} dev={crypt_device} luksuuid={uuid} dm={name}")
+
+        vmdb.progress(
+            f"LUKS: name={crypt_name} dev={crypt_dev} luksuuid={uuid} dm={crypt_name}"
+        )
         vmdb.progress(f"LUKS: {state.tags._tags}")
-        vmdb.progress("remembering LUKS device {} as {}".format(crypt_device, name))
+        vmdb.progress("remembering LUKS device {} as {}".format(crypt_dev, crypt_name))
+
+        if rmtmp:
+            os.remove(key_file)
+
+    def _write_temp(self, passord):
+        fd, filename = tempfile.mkstemp()
+        os.close(fd)
+        open(filename, "w").write(password)
+        return filename
+
+    def teardown(self, values, settings, state):
+        x = state.tmp_key_file
+        if x is not None and os.path.exists(x):
+            os.remove(x)
+
+        crypt_name = values["name"]
 
-        shutil.rmtree(tmp)
+        crypt_dev = "/dev/mapper/{}".format(crypt_name)
+        vmdb.runcmd(["cryptsetup", "close", crypt_dev])