feat: add auth_method and partition information

This commit is contained in:
Adrián Jesús Peña Rodríguez committed 2025-02-19 11:29:42 +01:00
1 parent eb0f6addb4
commit 772667c8cc
2 files changed
+32 -7

No files matched your search

+26 -3
View File
@@ -1,5 +1,5 @@
from datetime import datetime
from typing import Optional, Union
from typing import Optional, Tuple, Union
from pydantic import BaseModel, Field, ValidationError
@@ -275,8 +275,29 @@ class Finding(BaseModel):
)
raise error
@staticmethod
def _get_auth_method_and_partition(provider: Provider) -> Tuple[str, str]:
"""
Extract the authentication method and partition information based on the provider type.
"""
auth_method = ""
partition = ""
if provider.type == "aws":
auth_method = f"profile: {get_nested_attribute(provider, 'identity.profile')}"
partition = get_nested_attribute(provider, "identity.partition")
elif provider.type == "azure":
auth_method = f"{provider.identity.identity_type}: {provider.identity.identity_id}"
partition = get_nested_attribute(provider, "region_config.name")
elif provider.type == "gcp":
auth_method = f"Principal: {get_nested_attribute(provider, 'identity.profile')}"
elif provider.type == "kubernetes":
auth_method = "in-cluster" if provider.identity.context == "In-Cluster" else "kubeconfig"
elif provider.type == "microsoft365":
auth_method = f"{provider.identity.identity_type}: {provider.identity.identity_id}"
return auth_method, partition
@classmethod
def transform_api_finding(cls, finding) -> "Finding":
def transform_api_finding(cls, finding, provider) -> "Finding":
"""
Transform a FindingModel instance into an API-friendly Finding object.
@@ -289,12 +310,14 @@ class Finding(BaseModel):
Args:
finding (API Finding): An API Finding instance containing data from the database.
provider (Provider): the provider object.
Returns:
Finding: A new Finding instance populated with data from the provided model.
"""
output_data = {}
output_data["auth_method"] = ""
output_data["auth_method"], output_data["partition"] = cls._get_auth_method_and_partition(provider)
output_data["timestamp"] = finding.inserted_at
output_data["account_uid"] = finding.scan.provider.uid
output_data["account_name"] = ""
+6 -4
View File
@@ -91,6 +91,7 @@ class DummyResources:
class DummyProvider:
def __init__(self, uid):
self.uid = uid
self.type = "aws"
class DummyScan:
@@ -571,10 +572,10 @@ class TestFinding:
dummy_finding.resources = resources
# Call the transform_api_finding classmethod
finding_obj = Finding.transform_api_finding(dummy_finding)
finding_obj = Finding.transform_api_finding(dummy_finding, provider)
# Fields directly set in transform_api_finding
assert finding_obj.auth_method == ""
assert finding_obj.auth_method == "profile: "
assert finding_obj.timestamp == inserted_at
assert finding_obj.account_uid == "account123"
assert finding_obj.account_name == ""
@@ -713,10 +714,11 @@ class TestFinding:
the function logs the error and re-raises the exception.
For example, if the metadata dict is missing required keys.
"""
provider = DummyProvider(uid="account123")
# Create a dummy API finding that is missing some required metadata
dummy_finding = DummyAPIFinding()
dummy_finding.inserted_at = 1234567890
dummy_finding.scan = DummyScan(provider=DummyProvider(uid="account123"))
dummy_finding.scan = DummyScan(provider=provider)
dummy_finding.uid = "finding-uid-invalid"
dummy_finding.status = "PASS"
dummy_finding.status_extended = "extended"
@@ -730,4 +732,4 @@ class TestFinding:
dummy_finding.resources = DummyResources(resource)
with pytest.raises(KeyError):
Finding.transform_api_finding(dummy_finding)
Finding.transform_api_finding(dummy_finding, provider)