fix: support sealing and unsealing secret in multiple process(more than one workers started by gunicorn)

This commit is contained in:
fxiang21
2024-04-15 16:47:29 +08:00
parent dc2f6ba957
commit 3b84687f89
3 changed files with 104 additions and 57 deletions

View File

@@ -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"]:

View File

@@ -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)

View File

@@ -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