1
- import typing
1
+ from typing import TYPE_CHECKING , Generator , Optional , Any
2
2
import arrow
3
3
from restfly .utils import dict_flatten , dict_merge
4
4
import uuid
5
5
6
- if typing . TYPE_CHECKING :
6
+ if TYPE_CHECKING :
7
7
from tenable .io .exports .iterator import ExportsIterator
8
8
from tenable .sc .analysis import AnalysisResultsIterator
9
9
10
10
11
- def tvm_asset_cleanup (* assets_iters : 'ExportsIterator' ) -> dict :
11
+ def tvm_asset_cleanup (* assets_iters : 'ExportsIterator'
12
+ ) -> Generator [Any , Any , Any ]:
12
13
"""
13
14
A simple wrapper to coalesce the multiple terminated asset states within
14
15
TVM.
@@ -26,8 +27,9 @@ def tvm_asset_cleanup(*assets_iters: 'ExportsIterator') -> dict:
26
27
27
28
def tvm_merged_data (assets_iter : 'ExportsIterator' ,
28
29
vulns_iter : 'ExportsIterator' ,
29
- asset_fields : list [str ] = None
30
- ) -> dict :
30
+ asset_fields : Optional [list [str ]] = None ,
31
+ close_accepted : bool = True ,
32
+ ) -> Generator [Any , Any , Any ]:
31
33
"""
32
34
Merges the asset and vulnerability finding data together into a single
33
35
object and adds in a computed finding id based on the following attributes:
@@ -86,11 +88,19 @@ def spf(value: str) -> str:
86
88
))
87
89
f ['integration_pid_updated' ] = pid
88
90
91
+ # If accepted risks shoudl be flagged as closed, then we will replace
92
+ # the state field with "fixed" if the risk was indeed accepted.
93
+ sevmod = f .get ('severity_modification_type' )
94
+ if close_accepted and sevmod == 'ACCEPTED' :
95
+ f ['state' ] = 'FIXED'
96
+
89
97
# Return the augmented finding to the caller.
90
98
yield f
91
99
92
100
93
- def tsc_merged_data (* vuln_iters : 'AnalysisResultsIterator' ) -> dict :
101
+ def tsc_merged_data (* vuln_iters : 'AnalysisResultsIterator' ,
102
+ close_accepted : bool = True ,
103
+ ) -> Generator [Any , Any , Any ]:
94
104
"""
95
105
Flattens and extends the vulnerability results returned from the
96
106
Security Center analysis API. The following fields are added to the
@@ -124,7 +134,9 @@ def tsc_merged_data(*vuln_iters: 'AnalysisResultsIterator') -> dict:
124
134
# If the hasBeenMitigated flag was flipped, then the finding isn't
125
135
# open, but is reopened. We want to confer state accurately so we
126
136
# will check that here.
127
- if f ['hasBeenMitigated' ] == '1' and state == 'open' :
137
+ if close_accepted and f ['acceptRisk' ] == '1' :
138
+ f ['integration_state' ] = 'fixed'
139
+ elif f ['hasBeenMitigated' ] == '1' and state == 'open' :
128
140
f ['integration_state' ] = 'reopened'
129
141
else :
130
142
f ['integration_state' ] = state
0 commit comments