Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 89 additions & 15 deletions infra/enforcement/account_keys.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import datetime
import logging
import sys
from torch import diff
import yaml
import argparse
import os
Expand Down Expand Up @@ -106,6 +107,53 @@
return username.split(":", 1)[1].strip().lower()
return username

def _get_user_managed_keys_from_iam(self, account_email: str) -> List[str]:
""""
Retrieves the list of user-managed keys for a given service account from IAM.

Args:
account_email (str): The email of the service account to retrieve keys for.

Returns:
List[str]: A list of key IDs for the user-managed keys associated with the service account
"""
request = types.ListServiceAccountKeysRequest()
request.name = f"projects/{self.project_id}/serviceAccounts/{account_email}"
request.key_types = [types.ListServiceAccountKeysRequest.KeyType.USER_MANAGED]

try:
response = self.service_account_client.list_service_account_keys(request=request)
return [key.name.split("/")[-1] for key in response.keys]
except Exception as e:
self.logger.error(f"Failed to retrieve keys for service account '{account_email}': {e}")
return []

def _get_verified_keys_from_secret_manager(self, secret_name: str) -> List[str]:
"""
Retrieves the list of verified keys for a given service account from Secret Manager.

Args:
secret_name (str): The name of the secret to retrieve keys for.

Returns:
List[str]: A list of key IDs for the verified keys associated with the service account.
"""
verified_keys = []
parent = self.secret_client.secret_path(self.project_id, secret_name)

try:
versions = self.secret_client.list_secret_versions(request={"parent": parent})
for version in versions:
if version.state.name == secretmanager.SecretVersion.State.ENABLED:
response = self.secret_client.access_secret_version(request={"name": version.name})
data_str = response.payload.data.decode("UTF-8")
key_id = data_str.split(":",1)[0]
verified_keys.append(key_id)
return verified_keys
except Exception as e:
self.logger.error(f"Failed to retrieve verified keys from Secret Manager for secret '{secret_name}': {e}")
return []

def _get_all_live_service_accounts(self) -> List[str]:
"""
Retrieves all service accounts that are currently active (not disabled) in the project.
Expand Down Expand Up @@ -259,21 +307,35 @@
self.logger.info(f"No service account keys found in the {self.service_account_keys_file}.")

compliance_issues = []
live_service_accounts = self._get_all_live_service_accounts()
managed_secrets = self._get_all_live_managed_secrets()

# Check that all service accounts that exist are declared
for service_account in self._get_all_live_service_accounts():
for service_account in live_service_accounts:
if self._denormalize_account_email(service_account) not in [account["account_id"] for account in file_service_accounts]:
msg = f"Service account '{service_account}' is not declared in the service account keys file."
compliance_issues.append(msg)
self.logger.warning(msg)
else:
iam_keys = self._get_user_managed_keys_from_iam(service_account)
if iam_keys:
secret_name = f"{self._denormalize_account_email(service_account)}-key"
legal_keys = []
if secret_name in managed_secrets:
legal_keys = self._get_verified_keys_from_secret_manager(secret_name)
unmanaged_keys = set(iam_keys) - set(legal_keys)
for unmanaged_key in unmanaged_keys:
msg = f"SECURITY ALERT: Unmanaged key '{unmanaged_key}' detected on account '{service_account}'. This key was created outside of Beam's service account management system. "
compliance_issues.append(msg)
self.logger.warning(msg)

managed_secrets = self._get_all_live_managed_secrets()
extracted_secrets = [f"{self._denormalize_account_email(account['account_id'])}-key" for account in file_service_accounts]

# Check for managed secrets that are not declared
for secret in managed_secrets:
if secret not in extracted_secrets:
msg = f"Managed secret '{secret}' is not declared in the service account keys file."
masked_secret = f"{secret[:4]}***{secret[-4:]}" if len(secret) >= 8 else "***"
msg = f"Managed secret '{masked_secret}' is not declared in the service account keys file."
compliance_issues.append(msg)
self.logger.warning(msg)

Expand Down Expand Up @@ -307,23 +369,34 @@
"""
if not self.sending_client:
raise ValueError("SendingClient is required for creating announcements")

diff = self.check_compliance()

if not diff:
self.logger.info("No compliance issues found, no announcement will be created.")
return
return

title = f"Account Keys Compliance Issue Detected"
body = f"Account keys for project {self.project_id} are not compliant with the defined policies on {self.service_account_keys_file}\n\n"
for issue in diff:
body += f"- {issue}\n"
unmanaged_keys_issues = [issue for issue in diff if "SECURITY ALERT" in issue]
general_issues = [issue for issue in diff if "SECURITY ALERT" not in issue]

announcement = f"Dear team,\n\nThis is an automated notification about compliance issues detected in the Account Keys policy for project {self.project_id}.\n\n"
announcement += f"We found {len(diff)} compliance issue(s) that need your attention.\n"
announcement += f"\nPlease check the GitHub issue for detailed information and take appropriate action to resolve these compliance violations."
if general_issues:
self.logger.info(f"Found {len(general_issues)} general compliance issues. Triggering announcement...")
title = f"Account Keys Compliance Issue Detected"
body = f"Account keys for project {self.project_id} are not compliant with the defined policies on {self.service_account_keys_file}\n\n"
for issue in general_issues:
body += f"- {issue}\n"

self.sending_client.create_announcement(title, body, recipient, announcement)
announcement = f"Dear team,\n\nThis is an automated notification about compliance issues detected in the Account Keys policy for project {self.project_id}.\n\n"
announcement += f"We found {len(general_issues)} compliance issue(s) that need your attention.\n"
announcement += f"\nPlease check the GitHub issue for detailed information and take appropriate action to resolve these compliance violations."

self.sending_client.create_announcement(title, body, recipient, announcement)
if unmanaged_keys_issues:
self.logger.info(f"Found {len(unmanaged_keys_issues)} unmanaged key security alerts. Dispatching to GitHub security issue...")
self.sending_client.report_unmanaged_keys(self.project_id, unmanaged_keys_issues)
else:
self.logger.info("No unmanaged key security alerts found, Checking if there are open security issues to auto-close...")
self.sending_client.resolve_unmanaged_keys()

def print_announcement(self, recipient: str) -> None:
"""
Expand Down Expand Up @@ -382,7 +455,8 @@
# Check for managed secrets that are not declared, if not, add them
for secret in managed_secrets:
if secret not in extracted_secrets:
self.logger.info(f"Managed secret '{secret}' is not declared in the service account keys file, adding it")
masked_secret = f"{secret[:4]}***{secret[-4:]}" if len(secret) >= 8 else "***"
self.logger.info(f"Managed secret '{masked_secret}' is not declared in the service account keys file, adding it")
file_service_accounts.append({
"account_id": secret.strip("-key"),
"display_name": self._normalize_account_email(secret.strip("-key")),
Expand Down Expand Up @@ -422,7 +496,7 @@

def config_process() -> Dict[str, str]:
with open(CONFIG_FILE, "r") as file:
config = yaml.safe_load(file)

Check failure

Code scanning / CodeQL

Clear-text logging of sensitive information High

This expression logs
sensitive data (secret)
as clear text.
This expression logs
sensitive data (secret)
as clear text.

if not config:
raise ValueError("Configuration file is empty or invalid.")
Expand Down Expand Up @@ -514,7 +588,7 @@
logger.error(f"Unknown action: {action}")
return 1
except Exception as e:
logger.error(f"Error executing action '{action}': {e}")
logger.exception(f"Error executing action '{action}': {e}")
return 1

return 0
Expand Down
88 changes: 85 additions & 3 deletions infra/enforcement/sending.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,10 +102,21 @@ def _get_open_issues(self, title: str) -> List[GitHubIssue]:
Args:
title (str): The title of the GitHub issue.
"""
endpoint = f"search/issues/?q=is:issue+repo:{self.github_repo}+in:title+{title}+is:open"
endpoint = f"search/issues?q=is:issue+repo:{self.github_repo}+in:title+{title}+is:open"
response = self._make_github_request("GET", endpoint)
issues = response.json().get('items', [])
return [GitHubIssue(**issue) for issue in issues]
parsed_issues = []
for issue in issues:
parsed_issues.append(GitHubIssue(
number=issue.get("number"),
title=issue.get("title"),
body=issue.get("body"),
state=issue.get("state"),
html_url=issue.get("html_url"),
created_at=issue.get("created_at"),
updated_at=issue.get("updated_at")
))
return parsed_issues

def create_issue(self, title: str, body: str) -> GitHubIssue:
"""
Expand All @@ -119,7 +130,16 @@ def create_issue(self, title: str, body: str) -> GitHubIssue:
payload = {"title": title, "body": body}
response = self._make_github_request("POST", endpoint, json=payload)
self.logger.info(f"Successfully created GitHub issue: {title}")
return GitHubIssue(**response.json())
data = response.json()
return GitHubIssue(
number=data.get("number"),
title=data.get("title"),
body=data.get("body"),
state=data.get("state"),
html_url=data.get("html_url"),
created_at=data.get("created_at"),
updated_at=data.get("updated_at")
)

def update_issue_body(self, issue_number: int, new_body: str) -> None:
"""
Expand All @@ -134,6 +154,68 @@ def update_issue_body(self, issue_number: int, new_body: str) -> None:
self._make_github_request("PATCH", endpoint, json=payload)
self.logger.info(f"Successfully updated body on GitHub issue: #{issue_number}")

def create_issue_comment(self, issue_number: int, comment_body: str) -> None:
"""
Adds a new comment to an existing GitHub issue in the specified repository.

Args:
issue_number (int): The number of the GitHub issue to comment on.
comment_body (str): The content of the comment to add to the GitHub issue.
"""
endpoint = f"repos/{self.github_repo}/issues/{issue_number}/comments"
payload = {"body": comment_body}
self._make_github_request("POST", endpoint, json=payload)
self.logger.info(f"Successfully added comment to GitHub issue: #{issue_number}")

def report_unmanaged_keys(self, project_id: str, compilance_issues: List[str]) -> None:
"""
Report compliance issues regarding unmanaged keys into a single GitHub issue.
Creates a new issue if none exists, otherwise appends a comment to the open one

Args:
project_id (str): The ID of the project associated with the unmanaged keys.
compilance_issues (List[str]): A list of compliance issues related to the unmanaged keys.
"""
if not compilance_issues:
self.logger.info("No compliance issues to report to Github.")
return

issue_title = "[SECURITY] Action Required: Unmanaged Service Account Keys Detected"
#markdown body
timestamp = __import__("datetime").datetime.now(__import__("datetime").timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
body = f"### Unmanaged Keys Audit Report ({timestamp})\n"
body += f"The following unauthorized or unmanaged keys were detected in `{project_id}`:\n\n"
for issue_text in compilance_issues:
body += f"- {issue_text}\n"

body += "\n*Please investigate and revoke these keys if they are not part of the official rotation system.*"

open_issues = self._get_open_issues(issue_title)
if open_issues:
target_issue = open_issues[0]
self.logger.info(f"Appending report to existing security issue #{target_issue.number}")
self.update_issue_body(target_issue.number, body)
else:
self.logger.info("Creating new security issue for unmanaged keys report.")
new_issue = self.create_issue(issue_title, body)
self.logger.info(f"Created new security issue : {new_issue.html_url}.")

def resolve_unmanaged_keys(self) -> None:
"""
Finds any open security issues regarding rogue keys and automatically closes them
if the infrastructure is now healthy.
"""
issue_title = "[SECURITY] Action Required: Unmanaged Service Account Keys Detected"
open_issues = self._get_open_issues(issue_title)
if open_issues:
target_issue = open_issues[0]
self.logger.info(f"All rogue keys resolved! Auto-closing issue #{target_issue.number}.")
self.create_issue_comment(target_issue.number, "All previously reported unmanaged keys have been resolved. Closing this issue.")
endpoint = f"repos/{self.github_repo}/issues/{target_issue.number}"
payload = {"state": "closed"}
self._make_github_request("PATCH", endpoint, json=payload)


def create_announcement(self, title: str, body: str, recipient: str, announcement: str) -> None:
"""
This method sends an email with an announcement. The email will point to a GitHub issue.
Expand Down
Loading