|
1 | 1 | import io
|
2 |
| - |
3 | 2 | import requests
|
4 | 3 |
|
| 4 | +from typing import Union |
| 5 | +from time import sleep |
| 6 | + |
5 | 7 | from ReversingLabs.SDK.helper import DEFAULT_USER_AGENT, WrongInputError
|
6 |
| -from ReversingLabs.SDK.ticloud import FileAnalysis, DynamicAnalysis, AdvancedSearch, FileDownload |
| 8 | +from ReversingLabs.SDK.ticloud import FileAnalysis, DynamicAnalysis, FileDownload |
7 | 9 | from ReversingLabs.SDK.a1000 import A1000
|
8 | 10 |
|
9 | 11 |
|
@@ -64,57 +66,6 @@ def enriched_file_analysis(self, sample_hash):
|
64 | 66 | return {}
|
65 | 67 |
|
66 | 68 |
|
67 |
| -class SpectraAssureScenarios(object): |
68 |
| - """A class for scenarios that include RL Spectra Assure and other RL services from the RLSDK.""" |
69 |
| - |
70 |
| - def __init__(self, spectra_assure_client, verify_certs=True): |
71 |
| - |
72 |
| - self._spectra_assure_client = spectra_assure_client |
73 |
| - self.verify_certs = verify_certs |
74 |
| - |
75 |
| - def a1000_upload_to_assure(self, a1000_host, a1000_token, file_hash, filename, project, package, version): |
76 |
| - a1000_client = A1000( |
77 |
| - host=a1000_host, |
78 |
| - token=a1000_token, |
79 |
| - user_agent=DEFAULT_USER_AGENT, |
80 |
| - verify=self.verify_certs |
81 |
| - ) |
82 |
| - |
83 |
| - file_content = a1000_client.download_sample(sample_hash=file_hash).content |
84 |
| - |
85 |
| - response = self._spectra_assure_client.submit_package( |
86 |
| - file=file_content, |
87 |
| - filename=filename, |
88 |
| - project=project, |
89 |
| - package=package, |
90 |
| - version=version |
91 |
| - ) |
92 |
| - |
93 |
| - return response |
94 |
| - |
95 |
| - def ticloud_upload_to_assure(self, ticloud_host, ticloud_username, ticloud_password, file_hash, filename, project, |
96 |
| - package, version): |
97 |
| - download_client = FileDownload( |
98 |
| - host=ticloud_host, |
99 |
| - username=ticloud_username, |
100 |
| - password=ticloud_password, |
101 |
| - user_agent=DEFAULT_USER_AGENT, |
102 |
| - verify=self.verify_certs |
103 |
| - ) |
104 |
| - |
105 |
| - file_content = download_client.download_sample(hash_input=file_hash).content |
106 |
| - |
107 |
| - response = self._spectra_assure_client.submit_package( |
108 |
| - file=file_content, |
109 |
| - filename=filename, |
110 |
| - project=project, |
111 |
| - package=package, |
112 |
| - version=version |
113 |
| - ) |
114 |
| - |
115 |
| - return response |
116 |
| - |
117 |
| - |
118 | 69 | class SpectraAssureClient(object):
|
119 | 70 | def __init__(self, host, token, organization, group, user_agent=DEFAULT_USER_AGENT, verify=True):
|
120 | 71 | self._host = host
|
@@ -178,5 +129,196 @@ def get_analysis_report(self, report_type, project, package, version):
|
178 | 129 | return response
|
179 | 130 |
|
180 | 131 |
|
| 132 | +class SpectraAssureScenarios(object): |
| 133 | + """A class for scenarios that include RL Spectra Assure and other RL services from the RLSDK.""" |
181 | 134 |
|
| 135 | + def __init__(self, spectra_assure_client: SpectraAssureClient, verify_certs=True): |
| 136 | + self._spectra_assure_client = spectra_assure_client |
| 137 | + self.verify_certs = verify_certs |
| 138 | + |
| 139 | + def __fetch_and_upload_to_assure(self, client: Union[A1000, FileDownload], file_hash, project, |
| 140 | + starting_version, max_version): |
| 141 | + file_content = client.download_sample(file_hash).content |
| 142 | + status_code = 409 |
| 143 | + |
| 144 | + starting_version = int(starting_version * 10) |
| 145 | + max_version = int(max_version * 10) |
| 146 | + |
| 147 | + current_version = starting_version |
| 148 | + |
| 149 | + while status_code == 409: |
| 150 | + if current_version > max_version: |
| 151 | + break |
| 152 | + |
| 153 | + response = self._spectra_assure_client.submit_package( |
| 154 | + file=file_content, |
| 155 | + filename=file_hash, |
| 156 | + project=project, |
| 157 | + package=f"package-{file_hash}", |
| 158 | + version=float(current_version / 10) |
| 159 | + ) |
| 160 | + |
| 161 | + status_code = response.status_code |
| 162 | + current_version += 1 |
| 163 | + |
| 164 | + return float((current_version - 1) / 10) |
| 165 | + |
| 166 | + def __get_report_from_assure(self, file_hash, report_type, project, current_version, max_retries=6): |
| 167 | + if max_retries > 30: |
| 168 | + raise WrongInputError("Maximum number of fetch retries exceeded.") |
| 169 | + |
| 170 | + not_done = "version that is not done with analysis" |
| 171 | + |
| 172 | + response = self._spectra_assure_client.get_analysis_report( |
| 173 | + report_type=report_type, |
| 174 | + project=project, |
| 175 | + package=f"package-{file_hash}", |
| 176 | + version=current_version |
| 177 | + ) |
| 178 | + |
| 179 | + retry_attempt = 1 |
| 180 | + |
| 181 | + while not_done in response.text and retry_attempt <= max_retries: |
| 182 | + sleep(3) |
| 183 | + |
| 184 | + response = self._spectra_assure_client.get_analysis_report( |
| 185 | + report_type=report_type, |
| 186 | + project=project, |
| 187 | + package=f"package-{file_hash}", |
| 188 | + version=current_version |
| 189 | + ) |
| 190 | + |
| 191 | + retry_attempt += 1 |
| 192 | + |
| 193 | + return response |
| 194 | + |
| 195 | + def a1000_upload_to_assure(self, a1000_host, a1000_token, hash_list, project, starting_version=1.0, max_version=12.0, |
| 196 | + get_analysis_report=False, report_type=None, max_retries=6) -> dict: |
| 197 | + """Fetches a list of samples defined in the hash_list from Spectra Analyze and |
| 198 | + submits them to Spectra Assure for analysis. |
| 199 | + Since Spectra Assure requires specifying the version of the sample (package), |
| 200 | + this method tries increasing the version of each sample from the list until it reaches |
| 201 | + a version that does not exist yet. |
| 202 | + You can specify the starting version and the maximum version to try with before backing off for |
| 203 | + the sample that is currently being uploaded. |
| 204 | + :param a1000_host: host of the desired Spectra Analyze instance, including the protocol prefix |
| 205 | + :type a1000_host: str |
| 206 | + :param a1000_token: authorization token of your Spectra Analyze account |
| 207 | + :type a1000_token: str |
| 208 | + :param hash_list: list containing the hashes of samples you want to upload to Spectra Assure |
| 209 | + :type hash_list: list[str] |
| 210 | + :param project: project name on Spectra Assure |
| 211 | + :type project: str |
| 212 | + :param starting_version: the starting version to try with for the current sample that is being uploaded |
| 213 | + :type starting_version: float |
| 214 | + :param max_version: the maximum version to try with before backing off for the current sample |
| 215 | + that is being uploaded |
| 216 | + :type max_version: float |
| 217 | + :param get_analysis_report: fetch the analysis report from Spectra Assure |
| 218 | + :type get_analysis_report: bool |
| 219 | + :param report_type: type of report to fetch |
| 220 | + :type report_type: str |
| 221 | + :param max_retries: maximum number of retries while trying to fetch the report before backing off |
| 222 | + :type max_retries: int |
| 223 | + :return: a dict of analysis reports |
| 224 | + :rtype: dict |
| 225 | + """ |
| 226 | + if max_version > 12: |
| 227 | + raise WrongInputError("The max_version can not be higher than 12.") |
| 228 | + |
| 229 | + a1000_client = A1000( |
| 230 | + host=a1000_host, |
| 231 | + token=a1000_token, |
| 232 | + user_agent=DEFAULT_USER_AGENT, |
| 233 | + verify=self.verify_certs |
| 234 | + ) |
| 235 | + |
| 236 | + reports = {} |
| 237 | + |
| 238 | + for file_hash in hash_list: |
| 239 | + current_version = self.__fetch_and_upload_to_assure( |
| 240 | + client=a1000_client, |
| 241 | + file_hash=file_hash, |
| 242 | + project=project, |
| 243 | + starting_version=starting_version, |
| 244 | + max_version=max_version |
| 245 | + ) |
| 246 | + |
| 247 | + if get_analysis_report: |
| 248 | + response = self.__get_report_from_assure( |
| 249 | + file_hash=file_hash, |
| 250 | + report_type=report_type, |
| 251 | + project=project, |
| 252 | + current_version=current_version, |
| 253 | + max_retries=max_retries |
| 254 | + ) |
| 255 | + |
| 256 | + reports[file_hash] = response.json() |
| 257 | + |
| 258 | + return reports |
| 259 | + |
| 260 | + def ticloud_upload_to_assure(self, ticloud_host, ticloud_username, ticloud_password, hash_list, project, |
| 261 | + starting_version=1, max_version=12, get_analysis_report=False, report_type=None, |
| 262 | + max_retries=6) -> dict: |
| 263 | + """Fetches a list of samples defined in the hash_list from Spectra Intelligence and |
| 264 | + submits them to Spectra Assure for analysis. |
| 265 | + Since Spectra Assure requires specifying the version of the sample (package), |
| 266 | + this method tries increasing the version of each sample from the list until it reaches |
| 267 | + a version that does not exist yet. |
| 268 | + You can specify the starting version and the maximum version to try with before backing off for |
| 269 | + the sample that is currently being uploaded. |
| 270 | + :param ticloud_host: Spectra Intelligence host name |
| 271 | + :type ticloud_host: str |
| 272 | + :param ticloud_username: your Spectra Intelligence username |
| 273 | + :type ticloud_username: str |
| 274 | + :param ticloud_password: your Spectra Intelligence password |
| 275 | + :type ticloud_password: str |
| 276 | + :param hash_list: list containing the hashes of samples you want to upload to Spectra Assure |
| 277 | + :type hash_list: list[str] |
| 278 | + :param project: project name on Spectra Assure |
| 279 | + :type project: str |
| 280 | + :param starting_version: the starting version to try with for the current sample that is being uploaded |
| 281 | + :type starting_version: float |
| 282 | + :param max_version: the maximum version to try with before backing off for the current sample |
| 283 | + that is being uploaded |
| 284 | + :type max_version: float |
| 285 | + :param get_analysis_report: fetch the analysis report from Spectra Assure |
| 286 | + :type get_analysis_report: bool |
| 287 | + :param report_type: type of report to fetch |
| 288 | + :type report_type: str |
| 289 | + :param max_retries: maximum number of retries while trying to fetch the report before backing off |
| 290 | + :type max_retries: int |
| 291 | + :return: a dict of analysis reports |
| 292 | + :rtype: dict |
| 293 | + """ |
| 294 | + download_client = FileDownload( |
| 295 | + host=ticloud_host, |
| 296 | + username=ticloud_username, |
| 297 | + password=ticloud_password, |
| 298 | + user_agent=DEFAULT_USER_AGENT, |
| 299 | + verify=self.verify_certs |
| 300 | + ) |
182 | 301 |
|
| 302 | + reports = {} |
| 303 | + |
| 304 | + for file_hash in hash_list: |
| 305 | + current_version = self.__fetch_and_upload_to_assure( |
| 306 | + client=download_client, |
| 307 | + file_hash=file_hash, |
| 308 | + project=project, |
| 309 | + starting_version=starting_version, |
| 310 | + max_version=max_version |
| 311 | + ) |
| 312 | + |
| 313 | + if get_analysis_report: |
| 314 | + response = self.__get_report_from_assure( |
| 315 | + file_hash=file_hash, |
| 316 | + report_type=report_type, |
| 317 | + project=project, |
| 318 | + current_version=current_version, |
| 319 | + max_retries=max_retries |
| 320 | + ) |
| 321 | + |
| 322 | + reports[file_hash] = response.json() |
| 323 | + |
| 324 | + return reports |
0 commit comments