Added the ability to map to the account the asset is tied to #44 #43

Author: SteveMcGrath

tenable_aws_sechub/finding.py

@@ -3,11 +3,13 @@
 AWS Security Hub finding.
 """
 
-from typing import Dict
+import logging
+from typing import Dict, List
 
 import arrow
 from restfly.utils import dict_clean, dict_flatten, trunc
 
+log = logging.getLogger('sechub.finding')
 SEV_MAP = {0: 0, 1: 3, 2: 5, 3: 7, 4: 10}
 STATE_MAP = {
     'OPEN': 'ACTIVE',
@@ -24,16 +26,25 @@ class Finding:
 
     region: str
     account_id: str
+    allowed_accounts: List[str] | None = None
     start_date: str
     map_to_asset_account: bool
 
     def __init__(
-        self, region: str, account_id: str, map_to_asset_account: bool = False
+        self,
+        region: str,
+        account_id: str,
+        map_to_asset_account: bool = False,
+        allowed_accounts: List[str] | None = None,
     ):
         self.region = region
         self.account_id = account_id
         self.map_to_asset_account = map_to_asset_account
         self.start_date = arrow.now().isoformat()
+        if allowed_accounts:
+            self.allowed_accounts = allowed_accounts
+        elif map_to_asset_account and not allowed_accounts:
+            self.allowed_accounts = [account_id]
 
     def check_required_params(self, vuln: Dict):
         """
@@ -64,6 +75,15 @@ def check_required_params(self, vuln: Dict):
                     f' were not set on asset {vuln["asset.uuid"]}'
                 )
             )
+        if (
+            self.allowed_accounts
+            and vuln.get('asset.aws_owner_id') not in self.allowed_accounts
+            and self.map_to_asset_account
+        ):
+            raise KeyError(
+                f'asset {vuln["asset.aws_owner_id"]}:{vuln["asset.uuid"]} is not within'
+                ' one of the allowed accounts.'
+            )
 
     def generate(self, vuln: Dict) -> Dict:
         """

tenable_aws_sechub/transform.py

@@ -60,6 +60,7 @@ def __init__(self, config: Dict):
             region=config.get('aws_region', self.aws.meta.region_name),
             account_id=str(config['aws_account_id']),
             map_to_asset_account=config.get('map_to_asset_account', False),
+            allowed_accounts=config.get('allowed_accounts'),
         )
         self._log = logging.getLogger('Tenb2SecHub')
 

tenable_aws_sechub/version.py

@@ -1 +1 @@
-version = '2.0.8'
+version = '2.1.0'

tests/test_finding.py

@@ -22,7 +22,7 @@ def test_finding_required_typerror(finding_local):
             '01234567-1234-abcd-0987-01234567890a'
         ),
     ):
-        resp = f.check_required_params(finding_local)
+        f.check_required_params(finding_local)
 
 
 def test_finding_generate_finding_keyerror(finding_local, finding_aws):
@@ -33,11 +33,20 @@ def test_finding_generate_finding_keyerror(finding_local, finding_aws):
     """
     f = Finding('', '')
     with pytest.raises(KeyError):
-        resp = f.generate(finding_local)
+        f.generate(finding_local)
 
     del finding_aws['plugin.type']
     with pytest.raises(KeyError, match='plugin.type'):
-        resp = f.generate(finding_aws)
+        f.generate(finding_aws)
+
+
+def test_finding_account_restrictions(finding_aws):
+    f = Finding('', '600832220000', True)
+    f.generate(finding_aws)
+
+    f = Finding('', 'ACCOUNT-ID', True)
+    with pytest.raises(KeyError, match='s not within one of the allowed accounts'):
+        f.generate(finding_aws)
 
 
 def test_finding_generate_finding_success(finding_aws):
@@ -56,7 +65,7 @@ def test_finding_generate_finding_success(finding_aws):
     assert resp['LastObservedAt'] == '2018-12-14T12:07:38.155Z'
     assert (
         resp['ProductArn']
-        == 'arn:aws:securityhub:AWS-REGION-1::product/tenable/tenable-io'
+        == 'arn:aws:securityhub:AWS-REGION-1::product/tenable/vulnerability-management'
     )
     assert resp['AwsAccountId'] == 'ACCOUNT-ID'
     assert resp['GeneratorId'] == 'tenable-plugin-106875'
@@ -82,7 +91,7 @@ def test_finding_generate_finding_success(finding_aws):
     assert resp['Resources'][0]['Details']['AwsEc2Instance']['IpV4Addresses'] == [
         '192.168.101.249'
     ]
-    assert resp['Resources'][0]['Details']['AwsEc2Instance']['IpV6Addresses'] == []
+    assert 'IpV6Addresses' not in resp['Resources'][0]['Details']['AwsEc2Instance']
     assert resp['ProductFields']['CVE'] == ''
     assert resp['ProductFields']['Plugin Family'] == 'Debian Local Security Checks'
     assert resp['ProductFields']['Type'] == 'local'