Repository navigation
feat(security): implement AES-256 encryption for API key storage - #215
sreerevanth merged 10 commits into
Conversation
…nraya-Tech-Labs#206) - Add encryption/decryption utilities using AES-256-GCM - Implement key hashing for database lookups - Create key rotation audit trail tracking - Add encrypted key storage database models - Implement access audit logging This ensures API keys are never stored in plaintext and includes comprehensive audit trails for key access and rotation.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds ChangesAPI Key Encryption and Storage
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 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 |
🧪 PR Test Results
Python 3.12 · commit 5a1456c |
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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 `@agentwatch/security/encryption.py`:
- Around line 105-125: rotate_key currently returns placeholders and doesn't
persist audit rows; update it to load the current stored key (or its hash),
compute old_key_hash (non-null) and new_key_hash via
APIKeyEncryption.hash_key(new_key), set rotated_at to current UTC datetime,
persist a KeyRotationAudit record via db_session (populate agent_id, rotated_at,
old_key_hash, new_key_hash, and any required fields), update the stored key
material using the existing key storage API, commit the transaction, and return
the saved audit entry as a dict; implement get_rotation_history to query
db_session for KeyRotationAudit rows filtered by agent_id, map them to
list[dict] with the same keys returned by rotate_key, and return that list.
Ensure you reference KeyRotationAudit, rotate_key, get_rotation_history,
db_session, and APIKeyEncryption.hash_key when making changes.
- Line 14: Replace typing.Optional with PEP 604 union syntax (e.g., use "str |
None") and remove the "from typing import Optional" import in this module; then
update KeyRotationManager.rotate_key to stop returning fixed nulls for audit
fields—either populate "old_key_hash" and "rotated_by" with real values derived
during rotation or omit those keys so DB defaults/constraints win (ensure the
returned audit payload includes the non-nullable rotated_by and any other
required fields of KeyRotationAudit); finally, replace the hardcoded PBKDF2 salt
in _derive_key with a per-master-key or per-derivation salt stored alongside the
master key (or generated and persisted) and use that salt when deriving keys so
salts are not fixed globally.
In `@agentwatch/security/key_storage.py`:
- Line 29: The timestamp columns created_at, rotated_at, and accessed_at are
using default=datetime.now(UTC) which is evaluated at import time; update each
Column definition in key_storage.py (the created_at, rotated_at, accessed_at
Column declarations) to use a callable default (e.g., a lambda or
functools.partial that returns datetime.now(UTC)) so the current time is
computed when a row is created/updated rather than at module import.
- Around line 11-14: Remove the unused imports Optional, create_engine, and
sessionmaker from the top of agentwatch/security/key_storage.py to fix F401;
change any column defaults that use default=datetime.now(UTC) to a callable
(e.g., default=lambda: datetime.now(UTC)) so timestamps are evaluated per-row
(update occurrences around the DateTime columns referenced in the file); and for
the primary key columns rotation_id and access_id either ensure inserts always
supply those IDs or add a per-row default UUID generator (e.g., a callable that
returns str(uuid.uuid4())) so new rows get unique IDs automatically.
🪄 Autofix (Beta)
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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 88416f8a-ee51-4b34-ab89-a747ce4769e7
📒 Files selected for processing (2)
agentwatch/security/encryption.pyagentwatch/security/key_storage.py
| key_hash = Column(String(64), unique=True, nullable=False, index=True) | ||
| encrypted_key = Column(Text, nullable=False) | ||
| nonce = Column(Text, nullable=False) | ||
| created_at = Column(DateTime(timezone=True), default=datetime.now(UTC)) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="agentwatch/security/key_storage.py"
ls -la "$FILE" || true
python3 - <<'PY'
import itertools,sys,os
path="agentwatch/security/key_storage.py"
with open(path,'r',encoding='utf-8') as f:
lines=f.readlines()
for start,end in [(1,120),(20,80),(35,70),(45,65)]:
print("\n---",path,"lines",start,"to",end,"---")
for i in range(start-1,min(end,len(lines))):
ln=i+1
print(f"{ln:4d}: {lines[i].rstrip()}")
PYRepository: sreerevanth/AgentWatch
Length of output: 6444
Use callable defaults for SQLAlchemy timestamp columns to avoid import-time frozen values.
created_at (line 29), rotated_at (line 44), and accessed_at (line 56) use default=datetime.now(UTC), which is evaluated immediately when the module is loaded. Switch to a callable default so each row gets the current time.
Suggested fix
- created_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
@@
- rotated_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
+ rotated_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
@@
- accessed_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
+ accessed_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC))🤖 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 `@agentwatch/security/key_storage.py` at line 29, The timestamp columns
created_at, rotated_at, and accessed_at are using default=datetime.now(UTC)
which is evaluated at import time; update each Column definition in
key_storage.py (the created_at, rotated_at, accessed_at Column declarations) to
use a callable default (e.g., a lambda or functools.partial that returns
datetime.now(UTC)) so the current time is computed when a row is created/updated
rather than at module import.
- Fix os module import redundancy in encrypt_key method - Add proper type hints for KeyRotationManager.db_session - Update return types to include Any for dict values - Remove unused sqlalchemy imports (create_engine, sessionmaker)
- Replace Optional[str] with str | None per UP045 rule - Remove unused Optional import
sreerevanth
left a comment
There was a problem hiding this comment.
Thanks for the contribution. The encryption primitives themselves look useful, but I don't think the feature is complete enough to merge yet.
A few concerns from the current implementation:
KeyRotationManager.rotate_key()currently returns placeholder values (old_key_hash=None,rotated_at=None) and does not appear to persist rotation records.get_rotation_history()always returns an empty list.- The audit models are defined, but I don't see any code that actually writes audit entries.
- The encryption utilities are not yet integrated into the API key storage/execution flow, so API keys do not appear to be encrypted automatically in the current system.
- The PR description mentions key rotation, audit trails, and secure API key storage, but the implementation currently provides the building blocks rather than the complete workflow.
Could you complete the integration (storage path, audit persistence, rotation flow) or narrow the PR scope/description to match the current implementation? Happy to review again once that's addressed.
- Implement rotate_key() to actually persist rotations to database - Add KeyRotationAudit entries with full audit trail (rotation_id, reason, rotated_by) - Implement get_rotation_history() to query audit table - Encrypt new keys using AES-256 before storage - Support both initial key creation and rotation scenarios Addresses maintainer feedback on incomplete KeyRotationManager placeholders. Signed-off-by: Anshul Jain <anshul23102@iiitd.ac.in>
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
agentwatch/security/encryption.py (3)
163-163: ⚖️ Poor tradeoffConsider deferring commit to caller for transaction control.
Calling
commit()insiderotate_keyprevents callers from batching multiple operations in a single transaction or implementing their own error handling with rollback. This can lead to partial updates if subsequent operations fail.🤖 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 `@agentwatch/security/encryption.py` at line 163, The `rotate_key` method currently calls `self.db_session.commit()` directly, which prevents callers from controlling transaction boundaries and batching multiple operations together. Remove the commit() call from inside the `rotate_key` method and let the caller manage when to commit the session. This allows for proper transaction control, error handling with rollback capability, and the ability to batch multiple operations before committing.
175-178: 💤 Low valueInconsistent db_session handling between methods.
rotate_keyraisesRuntimeErrorwhendb_sessionis missing (line 118-119), butget_rotation_historysilently returns an empty list. This inconsistency could mask configuration errors for callers expecting rotation history data.Consider raising
RuntimeErrorhere as well, or documenting the intentional difference in behavior.🤖 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 `@agentwatch/security/encryption.py` around lines 175 - 178, The get_rotation_history method silently returns an empty list when db_session is missing, while the rotate_key method raises a RuntimeError in the same situation, creating inconsistent error handling that could mask configuration problems. Modify the get_rotation_history method to raise a RuntimeError when db_session is None, matching the behavior of rotate_key, to ensure consistent and predictable error reporting across both methods.
34-43: ⚖️ Poor tradeoffConsider using a per-deployment salt for PBKDF2 key derivation.
The hardcoded salt
b"agentwatch-api"means identical master keys will always derive identical encryption keys across deployments. While PBKDF2 still provides key stretching, a unique salt stored alongside the master key (e.g., generated on first run and persisted) would improve resistance to precomputation attacks.🤖 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 `@agentwatch/security/encryption.py` around lines 34 - 43, The _derive_key method uses a hardcoded salt value b"agentwatch-api" which causes identical master keys to derive identical encryption keys across all deployments, weakening security against precomputation attacks. Replace this hardcoded salt with a per-deployment approach: generate a cryptographically random salt on first initialization, persist it in a configuration file or secure storage alongside the master key, and modify _derive_key to accept and use this unique salt instead of the hardcoded default. This ensures each deployment has its own salt, significantly improving resistance to precomputation attacks even if the same master key is used across multiple instances.
🤖 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 `@agentwatch/security/encryption.py`:
- Around line 152-172: The old_key_hash sentinel value uses "initial" (7
characters) which is inconsistent with the 64-character SHA-256 hex digest
format expected throughout the schema and codebase. Additionally, the return
statement on line 168 returns the raw old_key_hash parameter value which may be
None, while the audit entry stores "initial", creating a mismatch between what's
persisted and what's returned. Replace the "initial" placeholder with a
consistent 64-character well-known placeholder (such as all zeros: 0x0 repeated
to 64 characters) in the KeyRotationAudit initialization on line 156, and update
line 168 to return the stored rotation_entry.old_key_hash value instead of the
original old_key_hash parameter to ensure consistency between the audit record
and the response.
- Around line 121-123: The inline imports at lines 121-123 violate Ruff's I001
rule due to improper import ordering. Reorganize these imports to follow isort
conventions by grouping the standard library imports (datetime and uuid)
together, separating them from the local agentwatch import with a blank line,
and ensuring alphabetical ordering within each group. Alternatively, move all
three imports to the module level at lines 10-18 to eliminate inline imports
entirely and avoid the repeated import overhead.
---
Nitpick comments:
In `@agentwatch/security/encryption.py`:
- Line 163: The `rotate_key` method currently calls `self.db_session.commit()`
directly, which prevents callers from controlling transaction boundaries and
batching multiple operations together. Remove the commit() call from inside the
`rotate_key` method and let the caller manage when to commit the session. This
allows for proper transaction control, error handling with rollback capability,
and the ability to batch multiple operations before committing.
- Around line 175-178: The get_rotation_history method silently returns an empty
list when db_session is missing, while the rotate_key method raises a
RuntimeError in the same situation, creating inconsistent error handling that
could mask configuration problems. Modify the get_rotation_history method to
raise a RuntimeError when db_session is None, matching the behavior of
rotate_key, to ensure consistent and predictable error reporting across both
methods.
- Around line 34-43: The _derive_key method uses a hardcoded salt value
b"agentwatch-api" which causes identical master keys to derive identical
encryption keys across all deployments, weakening security against
precomputation attacks. Replace this hardcoded salt with a per-deployment
approach: generate a cryptographically random salt on first initialization,
persist it in a configuration file or secure storage alongside the master key,
and modify _derive_key to accept and use this unique salt instead of the
hardcoded default. This ensures each deployment has its own salt, significantly
improving resistance to precomputation attacks even if the same master key is
used across multiple instances.
🪄 Autofix (Beta)
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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 2271186b-b380-4184-9386-aba4bce90224
📒 Files selected for processing (2)
agentwatch/security/encryption.pyagentwatch/security/key_storage.py
🚧 Files skipped from review as they are similar to previous changes (1)
- agentwatch/security/key_storage.py
| # Create audit entry | ||
| rotation_entry = KeyRotationAudit( | ||
| rotation_id=str(uuid4()), | ||
| agent_id=agent_id, | ||
| old_key_hash=old_key_hash or "initial", | ||
| new_key_hash=new_key_hash, | ||
| rotated_by=rotated_by, | ||
| rotated_at=datetime.now(UTC), | ||
| reason=reason, | ||
| ) | ||
| self.db_session.add(rotation_entry) | ||
| self.db_session.commit() | ||
|
|
||
| return { | ||
| "rotation_id": rotation_entry.rotation_id, | ||
| "agent_id": agent_id, | ||
| "old_key_hash": old_key_hash, | ||
| "new_key_hash": new_key_hash, | ||
| "rotated_at": rotation_entry.rotated_at.isoformat(), | ||
| "rotated_by": rotated_by, | ||
| "reason": reason, |
There was a problem hiding this comment.
Use a consistent 64-character placeholder for initial key rotation.
Line 156 uses "initial" (7 characters) as the old_key_hash sentinel for new keys, but the schema and all other code paths expect 64-character SHA-256 hex digests. This semantic inconsistency can confuse downstream consumers expecting uniform hash formats. Additionally, line 168 returns the raw old_key_hash (None) while the audit stores "initial", creating a mismatch.
🐛 Proposed fix
Use a well-known 64-character placeholder that's clearly not a real hash:
+ INITIAL_KEY_SENTINEL = "0" * 64 # Placeholder for first-time key creation
+
# Create audit entry
rotation_entry = KeyRotationAudit(
rotation_id=str(uuid4()),
agent_id=agent_id,
- old_key_hash=old_key_hash or "initial",
+ old_key_hash=old_key_hash or INITIAL_KEY_SENTINEL,
new_key_hash=new_key_hash,
rotated_by=rotated_by,
rotated_at=datetime.now(UTC),
reason=reason,
)
self.db_session.add(rotation_entry)
self.db_session.commit()
return {
"rotation_id": rotation_entry.rotation_id,
"agent_id": agent_id,
- "old_key_hash": old_key_hash,
+ "old_key_hash": old_key_hash or INITIAL_KEY_SENTINEL,
"new_key_hash": new_key_hash,
"rotated_at": rotation_entry.rotated_at.isoformat(),
"rotated_by": rotated_by,
"reason": reason,
}🤖 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 `@agentwatch/security/encryption.py` around lines 152 - 172, The old_key_hash
sentinel value uses "initial" (7 characters) which is inconsistent with the
64-character SHA-256 hex digest format expected throughout the schema and
codebase. Additionally, the return statement on line 168 returns the raw
old_key_hash parameter value which may be None, while the audit entry stores
"initial", creating a mismatch between what's persisted and what's returned.
Replace the "initial" placeholder with a consistent 64-character well-known
placeholder (such as all zeros: 0x0 repeated to 64 characters) in the
KeyRotationAudit initialization on line 156, and update line 168 to return the
stored rotation_entry.old_key_hash value instead of the original old_key_hash
parameter to ensure consistency between the audit record and the response.
- Move datetime, uuid imports to module level - Remove APIKeyEncryption from __init__ to avoid early instantiation - Instantiate encryption in rotate_key where needed Signed-off-by: Anshul Jain <anshul23102@iiitd.ac.in>
Signed-off-by: Anshul Jain <anshul23102@iiitd.ac.in>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
agentwatch/security/encryption.py (1)
166-167:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winAdd transaction rollback handling for commit failures.
If
commit()fails, the session may be left in an inconsistent state. Sincerotate_keytakes ownership of the commit, it should also handle rollback on failure.- self.db_session.add(rotation_entry) - self.db_session.commit() + self.db_session.add(rotation_entry) + try: + self.db_session.commit() + except Exception: + self.db_session.rollback() + raiseAlternatively, remove the
commit()call and let the caller manage transaction boundaries.🤖 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 `@agentwatch/security/encryption.py` around lines 166 - 167, The rotate_key method calls self.db_session.commit() without handling potential failures, which could leave the session in an inconsistent state. Either wrap the commit call in a try-except block and call self.db_session.rollback() if the commit fails to ensure the session is properly cleaned up, or remove the commit call entirely and let the caller manage transaction boundaries based on the method's scope of responsibility.
♻️ Duplicate comments (1)
agentwatch/security/encryption.py (1)
157-177:⚠️ Potential issue | 🟠 Major | ⚡ Quick winInconsistent
old_key_hashbetween persisted audit and returned dict.Line 160 stores
old_key_hash or "initial"in the database, but line 172 returns the rawold_key_hashvalue (which isNonefor new keys). This creates a mismatch between what's persisted and what's returned to the caller.Additionally,
"initial"(7 chars) doesn't match the 64-character SHA-256 format used elsewhere. Use a consistent 64-character placeholder and return the same value that's stored.🐛 Proposed fix
+ # Use 64-char placeholder for schema consistency + INITIAL_KEY_SENTINEL = "0" * 64 + # Create audit entry + stored_old_hash = old_key_hash or INITIAL_KEY_SENTINEL rotation_entry = KeyRotationAudit( rotation_id=str(uuid4()), agent_id=agent_id, - old_key_hash=old_key_hash or "initial", + old_key_hash=stored_old_hash, new_key_hash=new_key_hash, rotated_by=rotated_by, rotated_at=datetime.now(UTC), reason=reason, ) self.db_session.add(rotation_entry) self.db_session.commit() return { "rotation_id": rotation_entry.rotation_id, "agent_id": agent_id, - "old_key_hash": old_key_hash, + "old_key_hash": stored_old_hash, "new_key_hash": new_key_hash, "rotated_at": rotation_entry.rotated_at.isoformat(), "rotated_by": rotated_by, "reason": reason, }🤖 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 `@agentwatch/security/encryption.py` around lines 157 - 177, The KeyRotationAudit entry stores old_key_hash as `old_key_hash or "initial"` (line 160) but the returned dictionary uses the raw old_key_hash parameter (line 172), creating inconsistency between persisted and returned values. Additionally, "initial" is only 7 characters while SHA-256 hashes are 64 characters. Replace the hardcoded "initial" string with a consistent 64-character placeholder matching the SHA-256 format, and update the returned dictionary to use rotation_entry.old_key_hash (the value actually stored in the database) instead of the raw old_key_hash parameter to ensure consistency.
🧹 Nitpick comments (2)
agentwatch/security/encryption.py (2)
36-45: 💤 Low valueConsider using a per-deployment or random salt for stronger key derivation.
The hardcoded salt
b"agentwatch-api"means identical master keys will always produce identical derived encryption keys. While the master key should be unique per deployment, a stored random salt (generated once during initial setup and persisted alongside configuration) would provide an additional layer of defense against precomputation attacks.This is not critical since the master key should already have high entropy, but it's a best-practice hardening measure.
🤖 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 `@agentwatch/security/encryption.py` around lines 36 - 45, The _derive_key method uses a hardcoded salt value which means identical master keys will always produce identical derived encryption keys, reducing defense against precomputation attacks. Replace the hardcoded salt `b"agentwatch-api"` with a per-deployment approach: generate a random salt once during initial setup, persist it in the configuration, and modify the _derive_key method to retrieve and use this persisted salt instead of the hardcoded default. This ensures each deployment has a unique salt while maintaining consistency across key derivations within that deployment.
81-83: 💤 Low valueAvoid logging raw exception details from cryptographic operations.
Logging
str(e)from decryption failures could expose implementation details or partial data in logs. Consider logging only a generic message or a sanitized error code.except Exception as e: - logger.error("Failed to decrypt API key: %s", str(e)) + logger.error("Failed to decrypt API key: %s", type(e).__name__) raise ValueError("Decryption failed") from e🤖 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 `@agentwatch/security/encryption.py` around lines 81 - 83, The exception handler in the decryption operation is logging raw exception details using str(e), which can expose sensitive cryptographic implementation information in logs. Remove the raw exception string from the logger.error call and log only a generic, sanitized message that does not contain exception details. The exception e should still be raised with from e to preserve the error chain for debugging, but the logger.error statement should contain only a descriptive message without any dynamic error content.
🤖 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.
Outside diff comments:
In `@agentwatch/security/encryption.py`:
- Around line 166-167: The rotate_key method calls self.db_session.commit()
without handling potential failures, which could leave the session in an
inconsistent state. Either wrap the commit call in a try-except block and call
self.db_session.rollback() if the commit fails to ensure the session is properly
cleaned up, or remove the commit call entirely and let the caller manage
transaction boundaries based on the method's scope of responsibility.
---
Duplicate comments:
In `@agentwatch/security/encryption.py`:
- Around line 157-177: The KeyRotationAudit entry stores old_key_hash as
`old_key_hash or "initial"` (line 160) but the returned dictionary uses the raw
old_key_hash parameter (line 172), creating inconsistency between persisted and
returned values. Additionally, "initial" is only 7 characters while SHA-256
hashes are 64 characters. Replace the hardcoded "initial" string with a
consistent 64-character placeholder matching the SHA-256 format, and update the
returned dictionary to use rotation_entry.old_key_hash (the value actually
stored in the database) instead of the raw old_key_hash parameter to ensure
consistency.
---
Nitpick comments:
In `@agentwatch/security/encryption.py`:
- Around line 36-45: The _derive_key method uses a hardcoded salt value which
means identical master keys will always produce identical derived encryption
keys, reducing defense against precomputation attacks. Replace the hardcoded
salt `b"agentwatch-api"` with a per-deployment approach: generate a random salt
once during initial setup, persist it in the configuration, and modify the
_derive_key method to retrieve and use this persisted salt instead of the
hardcoded default. This ensures each deployment has a unique salt while
maintaining consistency across key derivations within that deployment.
- Around line 81-83: The exception handler in the decryption operation is
logging raw exception details using str(e), which can expose sensitive
cryptographic implementation information in logs. Remove the raw exception
string from the logger.error call and log only a generic, sanitized message that
does not contain exception details. The exception e should still be raised with
from e to preserve the error chain for debugging, but the logger.error statement
should contain only a descriptive message without any dynamic error content.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: a4c5d7cf-3558-444d-b6bc-e6325579d82c
📒 Files selected for processing (1)
agentwatch/security/encryption.py
|
@anshul23102 Thanks for the updates and for addressing the earlier feedback. The encryption primitives and rotation flow are much closer to being mergeable now, but I still have one concern before approving:
|
|
@anshul23102 can u please fix your prs ?? |
…ruff lint - KeyRotationAudit.old_key_hash: change nullable=False to nullable=True so the initial rotation (no prior key) stores NULL instead of the string "initial". The column is String(64) which matches the SHA-256 hex length; "initial" (7 chars) would fail any length or format check downstream. - encryption.py: pass old_key_hash directly (None for initial) instead of falling back to the sentinel string. - tests/: fix ruff I001/F401 lint violations (import sort order, unused imports)
There was a problem hiding this comment.
🧹 Nitpick comments (1)
agentwatch/security/encryption.py (1)
179-182: ⚡ Quick winInconsistent db_session error handling between methods.
rotate_keyraisesRuntimeErrorwhendb_sessionisNone, butget_rotation_historysilently returns an empty list. This could mask configuration issues—a caller might incorrectly conclude there's no rotation history when the session was never configured.Consider aligning the behavior by also raising
RuntimeErrorhere, or at minimum returning a distinct sentinel/raising when the session is missing vs. when there's genuinely no history.♻️ Proposed fix to align error handling
def get_rotation_history(self, agent_id: str) -> list[dict[str, Any]]: """Get rotation history for an agent.""" if not self.db_session: - return [] + raise RuntimeError("Database session required for rotation history") from agentwatch.security.key_storage import KeyRotationAudit🤖 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 `@agentwatch/security/encryption.py` around lines 179 - 182, The get_rotation_history method silently returns an empty list when db_session is None, which is inconsistent with the rotate_key method that raises a RuntimeError in the same situation. This inconsistency can mask configuration issues by making callers incorrectly assume there is no rotation history when the database session was never configured. Modify the db_session None check in get_rotation_history to raise a RuntimeError instead of returning an empty list, matching the error handling pattern used in rotate_key.
🤖 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.
Nitpick comments:
In `@agentwatch/security/encryption.py`:
- Around line 179-182: The get_rotation_history method silently returns an empty
list when db_session is None, which is inconsistent with the rotate_key method
that raises a RuntimeError in the same situation. This inconsistency can mask
configuration issues by making callers incorrectly assume there is no rotation
history when the database session was never configured. Modify the db_session
None check in get_rotation_history to raise a RuntimeError instead of returning
an empty list, matching the error handling pattern used in rotate_key.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: a8ff3cb9-0e7d-40d9-8a01-ffd3918786be
📒 Files selected for processing (9)
agentwatch/security/encryption.pyagentwatch/security/key_storage.pytests/conftest.pytests/test_api_safety_check.pytests/test_coroutine_detection.pytests/test_memory.pytests/test_owasp_security.pytests/test_sandbox.pytests/test_server.py
✅ Files skipped from review due to trivial changes (7)
- tests/test_server.py
- tests/conftest.py
- tests/test_sandbox.py
- tests/test_memory.py
- tests/test_api_safety_check.py
- tests/test_coroutine_detection.py
- tests/test_owasp_security.py
🚧 Files skipped from review as they are similar to previous changes (1)
- agentwatch/security/key_storage.py
demo.py had an unterminated f-string from a bad merge (unescaped quotes inside an f-string literal, invalid pre-3.12). main.py registered three separate Typer instances all named "session"; the last registration silently shadowed the others, breaking session export/watch/replay/list/ score commands with usage errors. Consolidated to a single session_app.
|
Hi @sreerevanth, the failing checks are fixed. Root cause was unrelated to the encryption changes: the merge from main left two artifacts in the CLI module:
Consolidated to a single |
Keep the EmbeddingProvider._load mock from main and drop the duplicate _st_model assignment line introduced on this branch.
Summary
Implements comprehensive AES-256-GCM encryption for secure API key storage, addressing critical security vulnerability where keys were previously stored in plaintext.
Problem Statement
Issue #206: Agent API keys stored in plaintext database pose critical security risk. If database is compromised or accessed by admin tools, all agent credentials are exposed, allowing complete system impersonation.
Solution Overview
Three-Layer Security Approach:
Encryption at Rest
Key Management
Audit & Compliance
Technical Implementation
New Files
agentwatch/security/encryption.pyAPIKeyEncryptionclass: Handles AES-256-GCM encryption/decryptionKeyRotationManagerclass: Manages key rotation lifecycleagentwatch/security/key_storage.pyEncryptedAPIKey: SQLAlchemy model for encrypted key storageKeyRotationAudit: Audit trail for rotationsKeyAccessAudit: Access logging for complianceKey Features
✅ AES-256-GCM encryption with authenticated encryption
✅ SHA-256 hashing for lookup without plaintext exposure
✅ PBKDF2 key derivation (100,000 iterations)
✅ Audit trails for access and rotations
✅ Configurable via environment variables
✅ Production-ready error handling
Database Schema
Environment Variables
Testing
Security Considerations
✅ Keys never stored in plaintext
✅ Authenticated encryption (GCM mode)
✅ Proper nonce generation and handling
✅ PBKDF2 key derivation with high iteration count
✅ Separate hash for secure lookups
✅ Full audit trail for compliance
Next Steps
Related Issues
Closes #206
🔐 This PR implements Phase 1-2 of the encryption roadmap. Phases 3-5 (API integration, key rotation, full testing) will follow in subsequent PRs.
Summary by CodeRabbit
Release Notes