From 3b84687f89d57779696b5e45edde50d4f9ba6204 Mon Sep 17 00:00:00 2001 From: fxiang21 Date: Mon, 15 Apr 2024 16:47:29 +0800 Subject: [PATCH] fix: support sealing and unsealing secret in multiple process(more than one workers started by gunicorn) --- cmdb-api/api/commands/click_cmdb.py | 4 +- cmdb-api/api/lib/secrets/inner.py | 155 ++++++++++++++++++---------- cmdb-api/api/lib/secrets/secrets.py | 2 +- 3 files changed, 104 insertions(+), 57 deletions(-) diff --git a/cmdb-api/api/commands/click_cmdb.py b/cmdb-api/api/commands/click_cmdb.py index fef71c4..950d84b 100644 --- a/cmdb-api/api/commands/click_cmdb.py +++ b/cmdb-api/api/commands/click_cmdb.py @@ -32,7 +32,7 @@ from api.lib.perm.acl.resource import ResourceCRUD from api.lib.perm.acl.resource import ResourceTypeCRUD from api.lib.perm.acl.role import RoleCRUD from api.lib.secrets.inner import KeyManage -from api.lib.secrets.inner import global_key_threshold +from api.lib.secrets.inner import global_key_threshold, secrets_shares from api.lib.secrets.secrets import InnerKVManger from api.models.acl import App from api.models.acl import ResourceType @@ -363,7 +363,7 @@ def cmdb_inner_secrets_unseal(address): for i in range(global_key_threshold): token = click.prompt(f'Enter unseal token {i + 1}', hide_input=True, confirmation_prompt=False) assert token is not None - resp = requests.post(address, headers={"Unseal-Token": token}) + resp = requests.post(address, headers={"Unseal-Token": token}, timeout=5) if resp.status_code == 200: KeyManage.print_response(resp.json()) if resp.json().get("status") in ["success", "skip"]: diff --git a/cmdb-api/api/lib/secrets/inner.py b/cmdb-api/api/lib/secrets/inner.py index 2d58763..4a79feb 100644 --- a/cmdb-api/api/lib/secrets/inner.py +++ b/cmdb-api/api/lib/secrets/inner.py @@ -1,22 +1,17 @@ +import json import os import secrets import sys -from base64 import b64decode, b64encode +import threading +from base64 import b64decode, b64encode from Cryptodome.Protocol.SecretSharing import Shamir -from colorama import Back -from colorama import Fore -from colorama import Style -from colorama import init as colorama_init +from colorama import Back, Fore, Style, init as colorama_init from cryptography.hazmat.backends import default_backend -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives import padding -from cryptography.hazmat.primitives.ciphers import Cipher -from cryptography.hazmat.primitives.ciphers import algorithms -from cryptography.hazmat.primitives.ciphers import modes +from cryptography.hazmat.primitives import hashes, padding +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes from cryptography.hazmat.primitives.ciphers.aead import AESGCM from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC - from flask import current_app global_iv_length = 16 @@ -28,9 +23,12 @@ backend_encrypt_key_name = "encrypt_key" backend_root_key_salt_name = "root_key_salt" backend_encrypt_key_salt_name = "encrypt_key_salt" backend_seal_key = "seal_status" + success = "success" seal_status = True +secrets_encrypt_key = "" +secrets_root_key = "" def string_to_bytes(value): if not value: @@ -47,6 +45,8 @@ def string_to_bytes(value): class Backend: def __init__(self, backend=None): self.backend = backend + # cache is a redis object + self.cache = backend.cache def get(self, key): return self.backend.get(key) @@ -76,11 +76,14 @@ class KeyManage: def init_app(self, app, backend=None): if (sys.argv[0].endswith("gunicorn") or (len(sys.argv) > 1 and sys.argv[1] in ("run", "cmdb-password-data-migrate"))): + + self.backend = backend + threading.Thread(target=self.watch_root_key, args=(app,)).start() + self.trigger = app.config.get("INNER_TRIGGER_TOKEN") if not self.trigger: return - self.backend = backend resp = self.auto_unseal() self.print_response(resp) @@ -147,36 +150,42 @@ class KeyManage: else: return "", True - def auth_root_secret(self, root_key): - msg, ok = self.is_valid_root_key(root_key) - if not ok: - return { - "message": msg, - "status": "failed" - } + def auth_root_secret(self, root_key, app): + with app.app_context(): + msg, ok = self.is_valid_root_key(root_key) + if not ok: + return { + "message": msg, + "status": "failed" + } - encrypt_key_aes = self.backend.get(backend_encrypt_key_name) - if not encrypt_key_aes: - return { - "message": "encrypt key is empty", - "status": "failed" - } + encrypt_key_aes = self.backend.get(backend_encrypt_key_name) + if not encrypt_key_aes: + return { + "message": "encrypt key is empty", + "status": "failed" + } - secrets_encrypt_key, ok = InnerCrypt.aes_decrypt(string_to_bytes(root_key), encrypt_key_aes) - if ok: - msg, ok = self.backend.update(backend_seal_key, "open") + secret_encrypt_key, ok = InnerCrypt.aes_decrypt(string_to_bytes(root_key), encrypt_key_aes) if ok: - current_app.config["secrets_encrypt_key"] = secrets_encrypt_key - current_app.config["secrets_root_key"] = root_key - # current_app.config["secrets_shares"] = [] - self.backend.set_shares(self.share_key, []) - return {"message": success, "status": success} - return {"message": msg, "status": "failed"} - else: - return { - "message": secrets_encrypt_key, - "status": "failed" - } + msg, ok = self.backend.update(backend_seal_key, "open") + if ok: + global secrets_encrypt_key, secrets_root_key + secrets_encrypt_key = secret_encrypt_key + secrets_root_key = root_key + self.backend.cache.set(self.share_key, json.dumps([])) + return {"message": success, "status": success} + return {"message": msg, "status": "failed"} + else: + return { + "message": secret_encrypt_key, + "status": "failed" + } + + def parse_shares(self, shares, app): + if len(shares) >= global_key_threshold: + recovered_secret = Shamir.combine(shares[:global_key_threshold], False) + return self.auth_root_secret(b64encode(recovered_secret), app) def unseal(self, key): if not self.is_seal(): @@ -188,16 +197,12 @@ class KeyManage: try: t = [i for i in b64decode(key)] v = (int("".join([chr(i) for i in t[-2:]])), bytes(t[:-2])) - # shares = current_app.config.get("secrets_shares", []) shares = self.backend.get_shares(self.share_key) if v not in shares: shares.append(v) - # current_app.config["secrets_shares"] = shares - self.backend.set_shares(self.share_key, shares) - + self.set_shares(shares) if len(shares) >= global_key_threshold: - recovered_secret = Shamir.combine(shares[:global_key_threshold], False) - return self.auth_root_secret(b64encode(recovered_secret)) + return self.parse_shares(shares, current_app) else: return { "message": "waiting for inputting other unseal key {0}/{1}".format(len(shares), @@ -257,8 +262,11 @@ class KeyManage: msg, ok = self.backend.add(backend_seal_key, "open") if not ok: return {"message": msg, "status": "failed"}, False - current_app.config["secrets_root_key"] = root_key - current_app.config["secrets_encrypt_key"] = encrypt_key + + global secrets_encrypt_key, secrets_root_key + secrets_encrypt_key = encrypt_key + secrets_root_key = root_key + self.print_token(shares, root_token=root_key) return {"message": "OK", @@ -281,7 +289,7 @@ class KeyManage: } # TODO elif len(self.trigger.strip()) == 24: - res = self.auth_root_secret(self.trigger.encode()) + res = self.auth_root_secret(self.trigger.encode(), current_app) if res.get("status") == success: return { "message": success, @@ -313,19 +321,26 @@ class KeyManage: "message": msg, "status": "failed", } - current_app.config["secrets_root_key"] = '' - current_app.config["secrets_encrypt_key"] = '' + self.clear() + self.backend.cache.publish(self.share_key, "clear") + return { "message": success, "status": success } + @staticmethod + def clear(): + global secrets_encrypt_key, secrets_root_key + secrets_encrypt_key = '' + secrets_root_key = '' + def is_seal(self): """ If there is no initialization or the root key is inconsistent, it is considered to be in a sealed state.. :return: """ - secrets_root_key = current_app.config.get("secrets_root_key") + # secrets_root_key = current_app.config.get("secrets_root_key") if not secrets_root_key: return True msg, ok = self.is_valid_root_key(secrets_root_key) @@ -366,22 +381,53 @@ class KeyManage: } print(status_colors.get(status, Fore.GREEN), message, Style.RESET_ALL) + def set_shares(self, values): + new_value = list() + for v in values: + new_value.append((v[0], b64encode(v[1]).decode("utf-8"))) + self.backend.cache.publish(self.share_key, json.dumps(new_value)) + self.backend.cache.set(self.share_key, json.dumps(new_value)) + + def watch_root_key(self, app): + pubsub = self.backend.cache.pubsub() + pubsub.subscribe(self.share_key) + + new_value = set() + for message in pubsub.listen(): + if message["type"] == "message": + if message["data"] == b"clear": + self.clear() + continue + try: + value = json.loads(message["data"].decode("utf-8")) + for v in value: + new_value.add((v[0], b64decode(v[1]))) + except Exception as e: + return [] + if len(new_value) >= global_key_threshold: + self.parse_shares(list(new_value), app) + new_value = set() + class InnerCrypt: def __init__(self): - secrets_encrypt_key = current_app.config.get("secrets_encrypt_key", "") - self.encrypt_key = b64decode(secrets_encrypt_key.encode("utf-8")) + self.encrypt_key = b64decode(secrets_encrypt_key) + #self.encrypt_key = b64decode(secrets_encrypt_key, "".encode("utf-8")) def encrypt(self, plaintext): """ encrypt method contain aes currently """ + if not self.encrypt_key: + return ValueError("secret is disabled, please seal firstly"), False return self.aes_encrypt(self.encrypt_key, plaintext) def decrypt(self, ciphertext): """ decrypt method contain aes currently """ + if not self.encrypt_key: + return ValueError("secret is disabled, please seal firstly"), False return self.aes_decrypt(self.encrypt_key, ciphertext) @classmethod @@ -398,6 +444,7 @@ class InnerCrypt: return b64encode(iv + ciphertext).decode("utf-8"), True except Exception as e: + return str(e), False @classmethod @@ -443,4 +490,4 @@ if __name__ == "__main__": t_ciphertext, status1 = c.encrypt(t_plaintext) print("Ciphertext:", t_ciphertext) decrypted_plaintext, status2 = c.decrypt(t_ciphertext) - print("Decrypted plaintext:", decrypted_plaintext) + print("Decrypted plaintext:", decrypted_plaintext) \ No newline at end of file diff --git a/cmdb-api/api/lib/secrets/secrets.py b/cmdb-api/api/lib/secrets/secrets.py index 89e0b01..01c71e9 100644 --- a/cmdb-api/api/lib/secrets/secrets.py +++ b/cmdb-api/api/lib/secrets/secrets.py @@ -5,9 +5,9 @@ from api.models.cmdb import InnerKV from api.extensions import rd - class InnerKVManger(object): def __init__(self): + self.cache = rd.r pass @classmethod