Get kms key from env (#4)

Co-authored-by: Melissa Greenbaum <melissa@Melissas-MacBook-Pro.local>

kms_ext/extension.py

@@ -7,6 +7,7 @@
 
 import boto3
 import dotenv
+import os
 import structlog
 from cryptography.hazmat.primitives import hashes, serialization
 from cryptography.hazmat.primitives.asymmetric import padding
@@ -18,7 +19,6 @@
 
 log = structlog.get_logger()
 
-
 class KMSCrypto:
  """Encrypt plaintext using KMS-compliant RSA algorithm."""
 
@@ -67,15 +67,13 @@ def invoke(self, *args: Any | None, **kwargs: Any) -> None:
 
  def encrypt(
  self,
- kms_key_id: str,
  public_key_path: Path,
  dotenv_path: Path = Path(".env"),
  output_path: Path = Path("secrets.yml"),
  ) -> Path:
  """Encrypt a given dotenv file with a given RSA public key (PEM file).
 
  Args:
- kms_key_id: The KMS key ID for the keypair being used.
  public_key_path: Path to RSA public key (PEM file).
  dotenv_path: Path to dotenv file (defaults to '.env')
  output_path: Path to output file (defaults to 'secrets.yml')
@@ -94,7 +92,7 @@ def encrypt(
  )
  env_vars.append(EnvVar(name=key, value=secret))
 
- secrets = SecretsFile(kms_key_id=kms_key_id, env=env_vars)
+ secrets = SecretsFile(env=env_vars)
 
  with open(output_path, "w") as secrets_file:
  secrets_file.write(secrets.yaml())
@@ -108,14 +106,19 @@ def decrypt(
  ) -> Path:
  client = boto3.client("kms")
 
+ try:
+ kms_key_id = os.environ["KMS_KEY_ID"]
+ except KeyError as ex:
+ raise Exception("The environment variable $KMS_KEY_ID must be set to decrypt") from ex
+
  with open(input_path) as ciphertext_file:
  secrets = SecretsFile.parse_raw(ciphertext_file.read())
 
  for env_var in secrets.env:
  ciphertext = base64.b64decode(env_var.value.ciphertext)
  response = client.decrypt(
  CiphertextBlob=ciphertext,
- KeyId=secrets.kms_key_id,
+ KeyId=kms_key_id,
  EncryptionAlgorithm=env_var.value.scheme, # "RSAES_OAEP_SHA_256" | "SYMMETRIC_DEFAULT" | "RSAES_OAEP_SHA_1" | "SM2PKE"
  )
  plaintext = response["Plaintext"].decode("utf-8")

kms_ext/main.py

@@ -42,14 +42,12 @@ def describe(
 
 @app.command()
 def encrypt(
- kms_key_id: str,
  public_key_path: Path,
  dotenv_path: Optional[Path] = typer.Option(Path(".env")),
  output_path: Optional[Path] = typer.Option(Path("secrets.yml")),
 ) -> None:
  """Encrypt a given dotenv file with a given RSA Public Key (PEM file)."""
  ext.encrypt(
- kms_key_id=kms_key_id,
  public_key_path=public_key_path,
  dotenv_path=dotenv_path,
  output_path=output_path,

kms_ext/secrets.py

@@ -15,5 +15,4 @@ class EnvVar(BaseModel):
 
 
 class SecretsFile(YamlModel):
- kms_key_id: str
  env: List[EnvVar]