feat(security): Implement KMS encryption for credentials - #1117
Open
vprashrex wants to merge 3 commits into
Open
feat(security): Implement KMS encryption for credentials#1117vprashrex wants to merge 3 commits into
vprashrex wants to merge 3 commits into
Conversation
… update related tests
📝 WalkthroughWalkthroughKMS-backed credentials now use AES-GCM envelope encryption with ChangesKMS envelope encryption
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant CredentialService
participant AWSKMS
participant AESGCM
CredentialService->>AWSKMS: GenerateDataKey
AWSKMS-->>CredentialService: Return wrapped and plaintext data keys
CredentialService->>AESGCM: Encrypt credentials locally
AESGCM-->>CredentialService: Return nonce and ciphertext
CredentialService->>CredentialService: Store kms.v2 envelope
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
github-actions
Bot
changed the title
OpenAPI changes ⚪ No API surface changesNote This PR does not modify the API contract.
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
vprashrex
commented
Aug 10, 2026
| def upgrade(): | ||
| bind = op.get_bind() | ||
| with Session(bind=bind) as session: | ||
| execute_credential_reencrypt(session=session) |
Collaborator
Author
There was a problem hiding this comment.
@kartpop This will invoke the sync function and execute execute_credential_reencrypt
| """ | ||
| try: | ||
| if encrypted_credentials.startswith(KMS_CIPHERTEXT_PREFIX): | ||
| if encrypted_credentials.startswith(KMS_ENVELOPE_PREFIX): |
Collaborator
Author
There was a problem hiding this comment.
This if-elif condition checks the format of the existing credential and then decrypts it using the corresponding decryption logic for that format.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/app/tests/core/test_security.py`:
- Around line 84-140: Update the new test methods in
backend/app/tests/core/test_security.py:84-140 and the test at
backend/app/tests/services/credentials/test_reencrypt.py:43 to annotate each
fixture parameter with its concrete fixture type and add -> None return
annotations. The site
backend/app/alembic/versions/076_reencrypt_credentials_envelope.py:25-33
requires no direct change.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 08063f7c-75f6-4688-a03d-f01dd8ea6c26
📒 Files selected for processing (5)
backend/app/alembic/versions/076_reencrypt_credentials_envelope.pybackend/app/core/security.pybackend/app/tests/core/test_security.pybackend/app/tests/services/credentials/test_reencrypt.pydocs/wiki/modules/platform.md
Comment on lines
+84
to
+140
| def test_kms_v2_envelope_roundtrip(self, kms_key): | ||
| creds = {"openai": {"api_key": "sk-envelope-123"}} | ||
|
|
||
| encrypted = encrypt_credentials(creds) | ||
|
|
||
| assert encrypted.startswith(KMS_ENVELOPE_PREFIX) | ||
| segments = encrypted[len(KMS_ENVELOPE_PREFIX) :].split(":") | ||
| assert len(segments) == 3 | ||
| for seg in segments: | ||
| base64.b64decode(seg) # each segment must be valid base64 | ||
| assert decrypt_credentials(encrypted) == creds | ||
|
|
||
| def test_kms_v2_large_payload_over_4096_bytes(self, kms_key): | ||
| # Direct KMS encrypt caps at 4096 bytes; envelope encryption has no such limit. | ||
| creds = { | ||
| "service_account": {"private_key": "k" * 5000, "client_email": "svc@x"} | ||
| } | ||
| assert len(json.dumps(creds).encode()) > 4096 | ||
|
|
||
| encrypted = encrypt_credentials(creds) | ||
|
|
||
| assert encrypted.startswith(KMS_ENVELOPE_PREFIX) | ||
| assert decrypt_credentials(encrypted) == creds | ||
|
|
||
| def test_v1_row_still_decrypts_with_v2_active(self, kms_key): | ||
| creds = {"api_key": "sk-v1-legacy"} | ||
| blob = security._kms_client.encrypt( | ||
| KeyId=kms_key, Plaintext=json.dumps(creds).encode() | ||
| )["CiphertextBlob"] | ||
| v1_ciphertext = KMS_CIPHERTEXT_PREFIX + base64.b64encode(blob).decode() | ||
|
|
||
| assert decrypt_credentials(v1_ciphertext) == creds | ||
|
|
||
| def test_new_writes_produce_v2(self, kms_key): | ||
| encrypted = encrypt_credentials({"api_key": "sk-new"}) | ||
|
|
||
| assert encrypted.startswith(KMS_ENVELOPE_PREFIX) | ||
| assert not encrypted.startswith(KMS_CIPHERTEXT_PREFIX) | ||
|
|
||
| def test_v2_tampered_ciphertext_raises(self, kms_key): | ||
| encrypted = encrypt_credentials({"api_key": "sk-tamper"}) | ||
| wrapped_b64, nonce_b64, ct_b64 = encrypted[len(KMS_ENVELOPE_PREFIX) :].split( | ||
| ":" | ||
| ) | ||
| ct = bytearray(base64.b64decode(ct_b64)) | ||
| ct[0] ^= 0xFF | ||
| tampered = ( | ||
| f"{KMS_ENVELOPE_PREFIX}{wrapped_b64}:{nonce_b64}:" | ||
| f"{base64.b64encode(bytes(ct)).decode()}" | ||
| ) | ||
|
|
||
| with pytest.raises(ValueError, match="Failed to decrypt credentials"): | ||
| decrypt_credentials(tampered) | ||
|
|
||
| def test_v2_wrong_segment_count_raises(self, kms_key): | ||
| with pytest.raises(ValueError, match="Failed to decrypt credentials"): | ||
| decrypt_credentials(f"{KMS_ENVELOPE_PREFIX}onlyoneseg") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Candidate files =="
git ls-files | rg 'backend/app/(tests/core/test_security.py|alembic/versions/076_reencrypt_credentials_envelope.py|tests/services/credentials/test_reencrypt.py)$' || true
echo
echo "== test_security.py relevant section =="
if [ -f backend/app/tests/core/test_security.py ]; then
nl -ba backend/app/tests/core/test_security.py | sed -n '75,145p'
fi
echo
echo "== alembic migration relevant section =="
if [ -f backend/app/alembic/versions/076_reencrypt_credentials_envelope.py ]; then
nl -ba backend/app/alembic/versions/076_reencrypt_credentials_envelope.py | sed -n '1,60p'
fi
echo
echo "== test_reencrypt.py relevant section =="
if [ -f backend/app/tests/services/credentials/test_reencrypt.py ]; then
nl -ba backend/app/tests/services/credentials/test_reencrypt.py | sed -n '30,55p'
fi
echo
echo "== Existing typing fixtures in tests/security/alembic to infer exact fixture names =="
rg -n 'def (kms_key|monkeypatch)|class|Session|encrypt_credentials|decrypt_credentials|KMS_ENVELOPE_PREFIX' backend/app/tests/core/test_security.py backend/app/alembic/versions/076_reencrypt_credentials_envelope.py backend/app/tests/services/credentials/test_reencrypt.py | sed -n '1,120p'Repository: ProjectTech4DevAI/kaapi-backend
Length of output: 437
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Candidate files =="
git ls-files | rg 'backend/app/(tests/core/test_security.py|alembic/versions/076_reencrypt_credentials_envelope.py|tests/services/credentials/test_reencrypt.py)$' || true
echo
echo "== test_security.py relevant section =="
if [ -f backend/app/tests/core/test_security.py ]; then
nl -ba backend/app/tests/core/test_security.py | sed -n '75,145p'
fi
echo
echo "== alembic migration relevant section =="
if [ -f backend/app/alembic/versions/076_reencrypt_credentials_envelope.py ]; then
nl -ba backend/app/alembic/versions/076_reencrypt_credentials_envelope.py | sed -n '1,60p'
fi
echo
echo "== test_reencrypt.py relevant section =="
if [ -f backend/app/tests/services/credentials/test_reencrypt.py ]; then
nl -ba backend/app/tests/services/credentials/test_reencrypt.py | sed -n '30,55p'
fi
echo
echo "== Existing typing fixtures in tests/security/alembic to infer exact fixture names =="
rg -n 'def (kms_key|monkeypatch)|class|Session|encrypt_credentials|decrypt_credentials|KMS_ENVELOPE_PREFIX' backend/app/tests/core/test_security.py backend/app/alembic/versions/076_reencrypt_credentials_envelope.py backend/app/tests/services/credentials/test_reencrypt.py | sed -n '1,120p'Repository: ProjectTech4DevAI/kaapi-backend
Length of output: 437
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
files = [
Path("backend/app/tests/core/test_security.py"),
Path("backend/app/alembic/versions/076_reencrypt_credentials_envelope.py"),
Path("backend/app/tests/services/credentials/test_reencrypt.py"),
]
for path in files:
print(f"\n== {path} lines 1-180 or relevant ==")
lines = path.read_text(encoding="utf-8").splitlines()
start, end = max(1, min(180, len(lines))), len(lines)
if path.name == "backend/app/tests/core/test_security.py":
start, end = max(1, 75), min(145, len(lines))
if path.name == "backend/app/alembic/versions/076_reencrypt_credentials_envelope.py":
start, end = max(1, 1), min(60, len(lines))
if path.name == "backend/app/tests/services/credentials/test_reencrypt.py":
start, end = max(1, 30), min(55, len(lines))
for i in range(start, end + 1):
print(f"{i:5}: {lines[i-1]}")
PYRepository: ProjectTech4DevAI/kaapi-backend
Length of output: 9435
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for path in [
Path("backend/app/tests/core/test_security.py"),
Path("backend/app/alembic/versions/076_reencrypt_credentials_envelope.py"),
Path("backend/app/tests/services/credentials/test_reencrypt.py"),
]:
print(f"\n== {path} lines 1-120 ==")
lines = path.read_text(encoding="utf-8").splitlines()
start, end = max(1, min(120, len(lines))), len(lines)
for i in range(start, end + 1):
print(f"{i:5}: {lines[i-1]}")
PY
echo
echo "== Read-only annotation extractor for target lines =="
python3 - <<'PY'
from pathlib import Path
import re
files = [
Path("backend/app/tests/core/test_security.py"),
Path("backend/app/alembic/versions/076_reencrypt_credentials_envelope.py"),
Path("backend/app/tests/services/credentials/test_reencrypt.py"),
]
def target_lines(path: Path):
text = path.read_text()
all_lines = text.splitlines()
matches = list(re.finditer(r'(def|async def)\s+([A-Za-z_]\w*)', text))
results = []
for m in matches:
lineno = text.count("\n", 0, m.start()) + 1
if (
path.name == "backend/app/tests/core/test_security.py" and lineno in range(84, 140 + 1)
) or (
path.name == "backend/app/alembic/versions/076_reencrypt_credentials_envelope.py" and lineno in range(25, 33 + 1)
) or (
path.name == "backend/app/tests/services/credentials/test_reencrypt.py" and lineno == 43
):
line = all_lines[lineno - 1]
has_self = bool(re.search(r'\bself\b', line))
has_self_annot = bool(re.search(r'\bself\s*:\s*[^=,\n:]+', line))
has_return_arrow = " -> " in line
results.append((lineno, line, has_self, has_self_annot, has_return_arrow))
return results
for path in files:
print(f"{path}")
for item in target_lines(path):
print(f" line {item[0]} self_annot={item[3]} return_annot={item[4]}")
print(f" {item[1]}")
PYRepository: ProjectTech4DevAI/kaapi-backend
Length of output: 12407
Add narrow type hints to the new KMS credential tests.
The test method signatures in backend/app/tests/core/test_security.py lines 84-140 still omit concrete annotations. Add each fixture parameter type and -> None; apply the same rule to the new test method at backend/app/tests/services/credentials/test_reencrypt.py#43.
🧰 Tools
🪛 ast-grep (0.45.0)
[info] 100-100: use jsonify instead of json.dumps for JSON output
Context: json.dumps(creds)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
[info] 110-110: use jsonify instead of json.dumps for JSON output
Context: json.dumps(creds)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
📍 Affects 3 files
backend/app/tests/core/test_security.py#L84-L140(this comment)backend/app/alembic/versions/076_reencrypt_credentials_envelope.py#L25-L33backend/app/tests/services/credentials/test_reencrypt.py#L43-L43
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@backend/app/tests/core/test_security.py` around lines 84 - 140, Update the
new test methods in backend/app/tests/core/test_security.py:84-140 and the test
at backend/app/tests/services/credentials/test_reencrypt.py:43 to annotate each
fixture parameter with its concrete fixture type and add -> None return
annotations. The site
backend/app/alembic/versions/076_reencrypt_credentials_envelope.py:25-33
requires no direct change.
Source: Coding guidelines
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/app/services/credentials/reencrypt.py`:
- Around line 47-49: Update the row handling in the re-encryption migration to
validate the complete v2 envelope shape, including its three required base64
segments, before skipping it; do not rely on startswith(KMS_ENVELOPE_PREFIX)
alone. Malformed prefixed values must continue to decrypt_credentials so they
raise and roll back the migration. Revise the nearby comment to explain that
skipping a valid v2 envelope keeps the backfill idempotent.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 4378ef1d-cc09-4b9f-8020-b851b011db78
📒 Files selected for processing (1)
backend/app/services/credentials/reencrypt.py
Comment on lines
+47
to
+49
| if row.credential.startswith(KMS_ENVELOPE_PREFIX): | ||
| # Already in the new envelope format; skip. | ||
| continue |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Validate the v2 envelope before skipping the row.
startswith(KMS_ENVELOPE_PREFIX) proves only the prefix. It does not prove that the value contains the three base64 segments required by decrypt_credentials in backend/app/core/security.py. A malformed value such as kms.v2: is skipped here, so the migration can commit while leaving an unreadable credential in storage. Validate the envelope shape before continue; otherwise let decrypt_credentials reject it and roll back the migration.
As per coding guidelines, comments must explain why rather than what; state that skipping a valid v2 envelope keeps the backfill idempotent.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@backend/app/services/credentials/reencrypt.py` around lines 47 - 49, Update
the row handling in the re-encryption migration to validate the complete v2
envelope shape, including its three required base64 segments, before skipping
it; do not rely on startswith(KMS_ENVELOPE_PREFIX) alone. Malformed prefixed
values must continue to decrypt_credentials so they raise and roll back the
migration. Revise the nearby comment to explain that skipping a valid v2
envelope keeps the backfill idempotent.
Source: Coding guidelines
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Issue
Closes #1114
Summary
Checklist
Before submitting a pull request, please ensure that you mark these task.
fastapi run --reload app/main.pyordocker compose upin the repository root and test.Notes
Please add here if any other information is required for the reviewer.
Summary by CodeRabbit
Security Improvements
Maintenance
Documentation