| Total Complexity | 44 |
| Total Lines | 301 |
| Duplicated Lines | 7.97 % |
| Changes | 0 | ||
Duplicate code is one of the most pungent code smells. A rule that is often used is to re-structure code once it is duplicated in three or more places.
Common duplication problems, and corresponding solutions are:
Complex classes like vtapi3.__main__ often do a lot of different things. To break such a class down, we need to identify a cohesive component within that class. A common approach to find such a component is to look for fields/methods that share the same prefixes, or suffixes.
Once you have determined the fields that belong together, you can apply the Extract Class refactoring. If the component makes sense as a sub-class, Extract Subclass is also a candidate, and is often faster.
| 1 | import os |
||
|
|
|||
| 2 | import json |
||
| 3 | import errno |
||
| 4 | import argparse |
||
| 5 | import vtapi3 |
||
| 6 | |||
| 7 | def get_environment_api_key(): |
||
| 8 | """Returns the value of the API key from environment variables. To work correctly, you need |
||
| 9 | to create an environment variable VT_API_KEY and write the current API key value to it. |
||
| 10 | |||
| 11 | Return: |
||
| 12 | The API key value from the environment variable VT_API_KEY. |
||
| 13 | |||
| 14 | Exception |
||
| 15 | VirusTotalAPIError(API key environment error): If there is no environment variable |
||
| 16 | VT_API_KEY. |
||
| 17 | |||
| 18 | """ |
||
| 19 | if 'VT_API_KEY' in os.environ: |
||
| 20 | result = os.environ['VT_API_KEY'] |
||
| 21 | else: |
||
| 22 | raise vtapi3.VirusTotalAPIError('API key environment error', errno.EINVAL) |
||
| 23 | return result |
||
| 24 | |||
| 25 | def get_file_id_to_analyse(file_path, api_key): |
||
| 26 | """Returns the file ID for further analysis on VirusTotal. |
||
| 27 | |||
| 28 | Args: |
||
| 29 | file_path: Path to the file for which you want to get the ID (str). |
||
| 30 | api_key: Access key to VirusTotal API functions (str). |
||
| 31 | |||
| 32 | Return: |
||
| 33 | The file ID if successful, or error message if not. |
||
| 34 | """ |
||
| 35 | vt_files = vtapi3.VirusTotalAPIFiles(api_key) |
||
| 36 | try: |
||
| 37 | result = vt_files.upload(file_path) |
||
| 38 | if vt_files.get_last_http_error() == vt_files.HTTP_OK: |
||
| 39 | result = json.loads(result) |
||
| 40 | result = 'File ID: ' + str(result['data']['id']) |
||
| 41 | else: |
||
| 42 | result = 'HTTP error ' + str(vt_files.get_last_http_error()) |
||
| 43 | return result |
||
| 44 | except vtapi3.VirusTotalAPIError as err: |
||
| 45 | return err |
||
| 46 | |||
| 47 | def get_file_scan_report(file_path, api_key): |
||
| 48 | """Returns an analysis report for a file uploaded to VirusTotal. |
||
| 49 | |||
| 50 | Args: |
||
| 51 | file_path: Path to the file to be analyzed on VirusTotal (str). |
||
| 52 | api_key: Access key to VirusTotal API functions (str). |
||
| 53 | |||
| 54 | Return: |
||
| 55 | The report on the results of the analysis if successful, or error message if not. |
||
| 56 | """ |
||
| 57 | vt_files = vtapi3.VirusTotalAPIFiles(api_key) |
||
| 58 | try: |
||
| 59 | result = vt_files.upload(file_path) |
||
| 60 | View Code Duplication | if vt_files.get_last_http_error() == vt_files.HTTP_OK: |
|
| 61 | result = json.loads(result) |
||
| 62 | file_id = str(result['data']['id']) |
||
| 63 | vt_analyses = vtapi3.VirusTotalAPIAnalyses(api_key) |
||
| 64 | result = vt_analyses.get_report(file_id) |
||
| 65 | if vt_analyses.get_last_http_error() == vt_analyses.HTTP_OK: |
||
| 66 | result = json.loads(result) |
||
| 67 | result = 'Analysis report:\n' + json.dumps(result, sort_keys=False, indent=4) |
||
| 68 | else: |
||
| 69 | result = 'HTTP error ' + str(vt_analyses.get_last_http_error()) |
||
| 70 | else: |
||
| 71 | result = 'HTTP error ' + str(vt_files.get_last_http_error()) |
||
| 72 | return result |
||
| 73 | except vtapi3.VirusTotalAPIError as err: |
||
| 74 | return err |
||
| 75 | |||
| 76 | def get_file_analyse_report(file_path, api_key): |
||
| 77 | """Returns an analysis report for a file in the VirusTotal database. |
||
| 78 | |||
| 79 | Args: |
||
| 80 | file_path: Path to the file for which you want to get the analysis report (str). |
||
| 81 | api_key: Access key to VirusTotal API functions (str). |
||
| 82 | |||
| 83 | Return: |
||
| 84 | The report on the results of the analysis if successful, or error message if not. |
||
| 85 | """ |
||
| 86 | try: |
||
| 87 | vt_files = vtapi3.VirusTotalAPIFiles(api_key) |
||
| 88 | file_id = vt_files.get_file_id(file_path) |
||
| 89 | result = vt_files.get_report(file_id) |
||
| 90 | if vt_files.get_last_http_error() == vt_files.HTTP_OK: |
||
| 91 | result = json.loads(result) |
||
| 92 | result = 'Analysis report:\n' + json.dumps(result, sort_keys=False, indent=4) |
||
| 93 | else: |
||
| 94 | result = 'HTTP error ' + str(vt_files.get_last_http_error()) |
||
| 95 | return result |
||
| 96 | except vtapi3.VirusTotalAPIError as err: |
||
| 97 | return err |
||
| 98 | |||
| 99 | def get_hash_report(hash_id, api_key): |
||
| 100 | """Returns an analysis report for a file in the VirusTotal database. |
||
| 101 | |||
| 102 | Args: |
||
| 103 | hash_id: Hash (SHA256, SHA1 or MD5) of the file for which you want to get an analysis |
||
| 104 | report (str). |
||
| 105 | api_key: Access key to VirusTotal API functions (str). |
||
| 106 | |||
| 107 | Return: |
||
| 108 | The report on the results of the analysis if successful, or error message if not. |
||
| 109 | """ |
||
| 110 | try: |
||
| 111 | vt_files = vtapi3.VirusTotalAPIFiles(api_key) |
||
| 112 | result = vt_files.get_report(hash_id) |
||
| 113 | if vt_files.get_last_http_error() == vt_files.HTTP_OK: |
||
| 114 | result = json.loads(result) |
||
| 115 | result = 'Analysis report:\n' + json.dumps(result, sort_keys=False, indent=4) |
||
| 116 | else: |
||
| 117 | result = 'HTTP error ' + str(vt_files.get_last_http_error()) |
||
| 118 | return result |
||
| 119 | except vtapi3.VirusTotalAPIError as err: |
||
| 120 | return err |
||
| 121 | |||
| 122 | def get_url_id_to_analyse(url, api_key): |
||
| 123 | """Returns the URL ID for further analysis on VirusTotal. |
||
| 124 | |||
| 125 | Args: |
||
| 126 | url: URL address for which you want to get an ID (str). |
||
| 127 | api_key: Access key to VirusTotal API functions (str). |
||
| 128 | |||
| 129 | Return: |
||
| 130 | The URL ID if successful, or error message if not. |
||
| 131 | """ |
||
| 132 | vt_urls = vtapi3.VirusTotalAPIUrls(api_key) |
||
| 133 | try: |
||
| 134 | result = vt_urls.upload(url) |
||
| 135 | if vt_urls.get_last_http_error() == vt_urls.HTTP_OK: |
||
| 136 | result = json.loads(result) |
||
| 137 | result = 'URL ID: ' + result['data']['id'] |
||
| 138 | else: |
||
| 139 | result = 'HTTP error ' + str(vt_urls.get_last_http_error()) |
||
| 140 | return result |
||
| 141 | except vtapi3.VirusTotalAPIError as err: |
||
| 142 | return err |
||
| 143 | |||
| 144 | def get_url_scan_report(url, api_key): |
||
| 145 | """Returns an analysis report for a URL uploaded to VirusTotal. |
||
| 146 | |||
| 147 | Args: |
||
| 148 | url: URL to be analyzed on VirusTotal (str). |
||
| 149 | api_key: Access key to VirusTotal API functions (str). |
||
| 150 | |||
| 151 | Return: |
||
| 152 | The report on the results of the analysis if successful, or error message if not. |
||
| 153 | """ |
||
| 154 | vt_urls = vtapi3.VirusTotalAPIUrls(api_key) |
||
| 155 | try: |
||
| 156 | result = vt_urls.upload(url) |
||
| 157 | View Code Duplication | if vt_urls.get_last_http_error() == vt_urls.HTTP_OK: |
|
| 158 | result = json.loads(result) |
||
| 159 | url_id = result['data']['id'] |
||
| 160 | vt_analyses = vtapi3.VirusTotalAPIAnalyses(api_key) |
||
| 161 | result = vt_analyses.get_report(url_id) |
||
| 162 | if vt_analyses.get_last_http_error() == vt_analyses.HTTP_OK: |
||
| 163 | result = json.loads(result) |
||
| 164 | result = 'Analysis report:\n' + json.dumps(result, sort_keys=False, indent=4) |
||
| 165 | else: |
||
| 166 | result = 'HTTP error ' + str(vt_analyses.get_last_http_error()) |
||
| 167 | else: |
||
| 168 | result = 'HTTP error ' + str(vt_urls.get_last_http_error()) |
||
| 169 | return result |
||
| 170 | except vtapi3.VirusTotalAPIError as err: |
||
| 171 | return err |
||
| 172 | |||
| 173 | def get_url_analyse_report(url, api_key): |
||
| 174 | """Returns an analysis report for a URL in the VirusTotal database. |
||
| 175 | |||
| 176 | Args: |
||
| 177 | url: URL for which you want to get the analysis report (str). |
||
| 178 | api_key: Access key to VirusTotal API functions (str). |
||
| 179 | |||
| 180 | Return: |
||
| 181 | The report on the results of the analysis if successful, or error message if not. |
||
| 182 | """ |
||
| 183 | vt_urls = vtapi3.VirusTotalAPIUrls(api_key) |
||
| 184 | try: |
||
| 185 | url_id = vt_urls.get_url_id_base64(url) |
||
| 186 | result = vt_urls.get_report(url_id) |
||
| 187 | if vt_urls.get_last_http_error() == vt_urls.HTTP_OK: |
||
| 188 | result = json.loads(result) |
||
| 189 | result = 'Analysis report:\n' + json.dumps(result, sort_keys=False, indent=4) |
||
| 190 | else: |
||
| 191 | result = 'HTTP error ' + str(vt_urls.get_last_http_error()) |
||
| 192 | return result |
||
| 193 | except vtapi3.VirusTotalAPIError as err: |
||
| 194 | return err |
||
| 195 | |||
| 196 | def get_ip_report(ip_address, api_key): |
||
| 197 | """Returns a report on the results of IP address analysis. |
||
| 198 | |||
| 199 | Args: |
||
| 200 | ip_address: IP address for which you want to get the analysis report (str). |
||
| 201 | api_key: Access key to VirusTotal API functions (str). |
||
| 202 | |||
| 203 | Return: |
||
| 204 | The report on the results of the analysis if successful, or error message if not. |
||
| 205 | """ |
||
| 206 | try: |
||
| 207 | vt_ip = vtapi3.VirusTotalAPIIPAddresses(api_key) |
||
| 208 | result = vt_ip.get_report(ip_address) |
||
| 209 | if vt_ip.get_last_http_error() == vt_ip.HTTP_OK: |
||
| 210 | result = json.loads(result) |
||
| 211 | result = 'Analysis report:\n' + json.dumps(result, sort_keys=False, indent=4) |
||
| 212 | else: |
||
| 213 | result = 'HTTP error ' + str(vt_ip.get_last_http_error()) |
||
| 214 | return result |
||
| 215 | except vtapi3.VirusTotalAPIError as err: |
||
| 216 | return err |
||
| 217 | |||
| 218 | def get_domain_report(domain, api_key): |
||
| 219 | """Returns a report on the results of domain analysis. |
||
| 220 | |||
| 221 | Args: |
||
| 222 | domain: Domain for which you want to get the analysis report (str). |
||
| 223 | api_key: Access key to VirusTotal API functions (str). |
||
| 224 | |||
| 225 | Return: |
||
| 226 | The report on the results of the analysis if successful, or error message if not. |
||
| 227 | """ |
||
| 228 | try: |
||
| 229 | vt_domain = vtapi3.VirusTotalAPIDomains(api_key) |
||
| 230 | result = vt_domain.get_report(domain) |
||
| 231 | if vt_domain.get_last_http_error() == vt_domain.HTTP_OK: |
||
| 232 | result = json.loads(result) |
||
| 233 | result = 'Analysis report:\n' + json.dumps(result, sort_keys=False, indent=4) |
||
| 234 | else: |
||
| 235 | result = 'HTTP error ' + str(vt_domain.get_last_http_error()) |
||
| 236 | return result |
||
| 237 | except vtapi3.VirusTotalAPIError as err: |
||
| 238 | return err |
||
| 239 | |||
| 240 | def create_cmd_parser(): |
||
| 241 | parser = argparse.ArgumentParser(prog='vtapi3') |
||
| 242 | parser.add_argument('resource', |
||
| 243 | help='Object that you want to analyse in VirusTotal (file, URL, IP address or domain)') |
||
| 244 | parser.add_argument('-fid', '--file-id', action='store_true', dest='file_id', |
||
| 245 | help='Getting the identifier of the file for further analysis') |
||
| 246 | parser.add_argument('-fsr', '--file-scan-report', action='store_true', dest='file_scan_report', |
||
| 247 | help='Getting a report on the results of scanning a file') |
||
| 248 | parser.add_argument('-far', '--file-analyse-report', action='store_true', dest='file_analyse_report', |
||
| 249 | help='Getting a report on the results of file analysis (enabled by default)') |
||
| 250 | parser.add_argument('-hr', '--hash-report', action='store_true', dest='hash_report', |
||
| 251 | help='Getting a report on the results of analyzing a file by its hash (SHA256, SHA1 or MD5)') |
||
| 252 | parser.add_argument('-uid', '--url-id', action='store_true', dest='url_id', |
||
| 253 | help='Getting the identifier of the URL for further analysis') |
||
| 254 | parser.add_argument('-usr', '--url-scan-report', action='store_true', dest='url_scan_report', |
||
| 255 | help='Getting a report on the results of scanning a URL') |
||
| 256 | parser.add_argument('-uar', '--url-analyse-report', action='store_true', dest='url_analyse_report', |
||
| 257 | help='Getting a report on the results of URL analysis') |
||
| 258 | parser.add_argument('-ipr', '--ip-report', action='store_true', dest='ip_report', |
||
| 259 | help='Getting a report on the results of IP address analysis') |
||
| 260 | parser.add_argument('-dr', '--domain-report', action='store_true', dest='domain_report', |
||
| 261 | help='Getting a report on the results of domain analysis') |
||
| 262 | return parser |
||
| 263 | |||
| 264 | def get_cmd_options(parser): |
||
| 265 | return parser.parse_args() |
||
| 266 | |||
| 267 | def main(options): |
||
| 268 | print('\nThe vtapi3 package. Implements the VirusTotal service API functions (3 versions).') |
||
| 269 | print('MIT Copyright (c) 2020, Evgeny Drobotun\n') |
||
| 270 | try: |
||
| 271 | api_key = get_environment_api_key() |
||
| 272 | if options.file_id: |
||
| 273 | result = get_file_id_to_analyse(options.resource, api_key) |
||
| 274 | elif options.file_scan_report: |
||
| 275 | result = get_file_scan_report(options.resource, api_key) |
||
| 276 | elif options.file_analyse_report: |
||
| 277 | result = get_file_analyse_report(options.resource, api_key) |
||
| 278 | elif options.hash_report: |
||
| 279 | result = get_hash_report(options.resource, api_key) |
||
| 280 | elif options.url_id: |
||
| 281 | result = get_url_id_to_analyse(options.resource, api_key) |
||
| 282 | elif options.url_scan_report: |
||
| 283 | result = get_url_scan_report(options.resource, api_key) |
||
| 284 | elif options.url_analyse_report: |
||
| 285 | result = get_url_analyse_report(options.resource, api_key) |
||
| 286 | elif options.ip_report: |
||
| 287 | result = get_ip_report(options.resource, api_key) |
||
| 288 | elif options.domain_report: |
||
| 289 | result = get_domain_report(options.resource, api_key) |
||
| 290 | else: |
||
| 291 | result = get_file_analyse_report(options.resource, api_key) |
||
| 292 | return result |
||
| 293 | except vtapi3.VirusTotalAPIError as err: |
||
| 294 | print(err) |
||
| 295 | print('\nTo work correctly, you need the VT_API_KEY environment variable') |
||
| 296 | print('with the current access key to the VirusTotal API functions.') |
||
| 297 | return err |
||
| 298 | |||
| 299 | if __name__ == '__main__': |
||
| 300 | print(main(get_cmd_options(create_cmd_parser()))) |
||
| 301 |