From 177156b81dc46a8fcaec4febcfbc4e5fc9af50b3 Mon Sep 17 00:00:00 2001 From: Jason Hollis Date: Tue, 12 Dec 2017 14:16:29 -0500 Subject: [PATCH] Made code PEP8 Compliant ( except for 'too complex [mccabe]' errors (SecurityCenter4.(query/asset_update) & SecurityCenter5.analysis --- securitycenter/base.py | 6 +- securitycenter/nessus.py | 16 +++-- securitycenter/pvs.py | 12 ++-- securitycenter/sc4.py | 104 ++++++++++++++++++------------ securitycenter/sc5.py | 134 ++++++++++++++++++++++++--------------- 5 files changed, 165 insertions(+), 107 deletions(-) diff --git a/securitycenter/base.py b/securitycenter/base.py index 6866a62..8138ac6 100644 --- a/securitycenter/base.py +++ b/securitycenter/base.py @@ -23,7 +23,8 @@ class BaseAPI(object): _session = None _timeout = (30, 300) - def __init__(self, host, port=443, ssl_verify=False, scheme='https', log=False, timeout=None): + def __init__(self, host, port=443, ssl_verify=False, scheme='https', + log=False, timeout=None): '''BaseAPI Initialization''' self._session = requests.Session() self._host = host @@ -48,7 +49,8 @@ def _reset_session(self): self._session = requests.Session() def _url(self, path): - return '%s://%s:%s/%s%s' % (self._scheme, self._host, self._port, self._pre, path) + return '%s://%s:%s/%s%s' % (self._scheme, self._host, self._port, + self._pre, path) def _builder(self, **kwargs): if 'headers' not in kwargs: diff --git a/securitycenter/nessus.py b/securitycenter/nessus.py index bf4ba0b..2b7cecc 100644 --- a/securitycenter/nessus.py +++ b/securitycenter/nessus.py @@ -1,4 +1,4 @@ -from .base import BaseAPI, APIError, logging +from .base import BaseAPI, APIError class Nessus(BaseAPI): @@ -6,9 +6,11 @@ class Nessus(BaseAPI): _secret = None managed = False enterprise = False - def __init__(self, host, port=8834, ssl_verify=False, scheme='https', log=False, timeout=None): + + def __init__(self, host, port=8834, ssl_verify=False, scheme='https', + log=False, timeout=None): BaseAPI.__init__(self, host, port, ssl_verify, scheme, log, timeout) - #try: + d = self.get('server/properties').json() try: if 'managed' in d: @@ -26,14 +28,16 @@ def __init__(self, host, port=8834, ssl_verify=False, scheme='https', log=False, def _builder(self, **kwargs): kwargs = BaseAPI._builder(self, **kwargs) if self._access and self._secret: - kwargs['headers']['X-APIKeys'] = 'accessKey=%s; secretKey=%s' % (self._access, self._secret) + kwargs['headers']['X-APIKeys'] = ('accessKey=%s; secretKey=%s' + % (self._access, self._secret)) elif self._token: kwargs['headers']['X-Cookie'] = 'token=%s' % self._token return kwargs def login(self, username=None, password=None, access=None, secret=None): if username and password: - resp = self.post('session', json={'username': username, 'password': password}) + resp = self.post('session', + json={'username': username, 'password': password}) if resp.status_code == 200: self._token = resp.json()['token'] else: @@ -43,5 +47,3 @@ def login(self, username=None, password=None, access=None, secret=None): self._secret = secret else: raise APIError(404, 'No Authentication Methods Found') - - diff --git a/securitycenter/pvs.py b/securitycenter/pvs.py index ef7d7fd..03540c8 100644 --- a/securitycenter/pvs.py +++ b/securitycenter/pvs.py @@ -1,8 +1,9 @@ -from .base import BaseAPI, APIError, logging +from .base import BaseAPI, APIError class PVS(BaseAPI): - def __init__(self, host, port=8835, ssl_verify=False, scheme='https', log=False, timeout=None): + def __init__(self, host, port=8835, ssl_verify=False, scheme='https', + log=False, timeout=None): BaseAPI.__init__(self, host, port, ssl_verify, scheme, log, timeout) def _builder(self, **kwargs): @@ -10,11 +11,8 @@ def _builder(self, **kwargs): return kwargs def login(self, username, password): - resp = self.post('login', data={ - 'login': username, - 'password': password, - 'json': 1, - }) + resp = self.post('login', data={'login': username, + 'password': password, 'json': 1}) if resp.status_code == 200: self._token = resp.json()['reply']['contents']['token'] else: diff --git a/securitycenter/sc4.py b/securitycenter/sc4.py index 2ea2e72..dee666d 100644 --- a/securitycenter/sc4.py +++ b/securitycenter/sc4.py @@ -1,13 +1,17 @@ from datetime import date, datetime, timedelta -import random, calendar, json, sys +import random +import calendar +import json +import sys from zipfile import ZipFile from .base import APIError, BaseAPI -if sys.version_info > (3,0): +if sys.version_info > (3, 0): from io import StringIO else: from StringIO import StringIO + class SecurityCenter4(BaseAPI): ''' Connects to the SecurityCenter API based on the parameters specified. @@ -33,7 +37,8 @@ class SecurityCenter4(BaseAPI): 'SSA', ] - def __init__(self, host, port=443, ssl_verify=False, scheme='https', log=False, timeout=None): + def __init__(self, host, port=443, ssl_verify=False, scheme='https', + log=False, timeout=None): BaseAPI.__init__(self, host, port, ssl_verify, scheme, log, timeout) self.system = self._system() self.version = self.system['version'] @@ -96,7 +101,7 @@ def _builder(self, **kwargs): if self._token: kwargs['data']['token'] = self._token - # Then the module, action, and input fields as required from the API. + # Then the module, action, and input fields as required from the API. if 'module' in kwargs: kwargs['data']['module'] = kwargs['module'] del kwargs['module'] @@ -107,7 +112,7 @@ def _builder(self, **kwargs): # The input data must be dumped as a string value in order to # be correctly interpreted. kwargs['data']['input'] = json.dumps(kwargs['input']) - del kwargs['input'] + del kwargs['input'] return kwargs def raw_query(self, module, action, data=None, dejson=True, **kwargs): @@ -119,7 +124,8 @@ def raw_query(self, module, action, data=None, dejson=True, **kwargs): data = self.post('', **kwargs) if dejson: if data.json()['error_code']: - raise APIError(data.json()['error_code'], data.json()['error_msg']) + raise APIError(data.json()['error_code'], + data.json()['error_msg']) return data.json()['response'] else: return data @@ -161,10 +167,15 @@ def query(self, tool, filters=None, source='cumulative', sort=None, """ - data = [] # This is the list that we will be returning back to - # the calling function once we complete. - payload = {} # The dataset that we will be sending to the API via the - # raw_query function. + # This is the list that we will be returning back to + # the calling function once we complete. + + data = [] + + # The dataset that we will be sending to the API via the + # raw_query function. + + payload = {} # A simple data dictionary to determine the module that we will be used stype = { @@ -177,7 +188,8 @@ def query(self, tool, filters=None, source='cumulative', sort=None, # When the source is "individual", scan and directory should be provided # as well, and will be set in the payload. - if source == "individual" and scan is not None and directory is not None: + if source == ("individual" and scan is not None + and directory is not None): # convert directory passed as datetime to string if necessary if isinstance(directory, date): directory = directory.strftime("%Y-%m-%d") @@ -215,17 +227,24 @@ def query(self, tool, filters=None, source='cumulative', sort=None, # Now that we have everything we need, a quick sanity check first to # make sure some idiot didn't give us a completely empty filterset to # work with. If they did for some reason, just return an empty list. - #if len(filters) < 1: + # if len(filters) < 1: # return [] # Everything is set, checks out, and is ready to go. Now we have to # start running through the query loop and actually pull everything # together. We know that we will - items = [] # This is the resultset. It'll be different every time - # we loop. to get things going however, we will just - # set it to an empty list. - count = 0 # A simple counter to track the total number of results - # that have been returned. + + # This is the resultset. It'll be different every time + # we loop. to get things going however, we will just + # set it to an empty list. + + items = [] + + # A simple counter to track the total number of results + # that have been returned. + + count = 0 + while len(items) == req_size or count == 0: # The API requires that we set an offset for the start and end of # the request, so we will add these to the payload here. @@ -263,7 +282,7 @@ def login(self, user, passwd): queries. """ data = self.raw_query('auth', 'login', - data={'username': user, 'password': passwd}) + data={'username': user, 'password': passwd}) self._token = data["token"] self._user = data @@ -318,7 +337,6 @@ def asset_update(self, asset_id, name=None, description=None, if asset['type'] == 'dnsname': payload['definedDNSNames'] = asset['definedDNSNames'] - # New we need to check to see if we actually got to pre-load the # payload. If we didnt, then there isn an existing Asset list and we # should error out. @@ -344,12 +362,14 @@ def asset_update(self, asset_id, name=None, description=None, for user in users: ulist.append({'id': int(user)}) payload['users'] = ulist - if payload['type'] == 'dynamic' and rules is not None and isinstance(rules, list): + if payload['type'] == ('dynamic' and rules is not None + and isinstance(rules, list)): payload['rules'] = rules - if payload['type'] == 'static' and ips is not None and isinstance(ips, list): + if payload['type'] == ('static' and ips is not None + and isinstance(ips, list)): payload['definedIPs'] = ','.join(ips) - if payload['type'] == 'dnsname' and dns is not None\ - and isinstance(dns, list): + if payload['type'] == ('dnsname' and dns is not None + and isinstance(dns, list)): payload['definedDNSNames'] = ','.join(dns) # And now that we have everything defined, we can go ahead and send @@ -401,7 +421,8 @@ def credential_update(self, cred_id, **options): payload['username'] = cred['username'] payload['publickey'] = cred['publickey'] payload['privatekey'] = cred['privatekey'] - payload['priviledgeEscalation'] = cred['priviledgeEscalation'] + payload['priviledgeEscalation'] = ( + cred['priviledgeEscalation']) payload['escalationUsername'] = cred['escalationUsername'] if cred['type'] == 'windows': @@ -520,9 +541,11 @@ def credential_add(self, name, cred_type, **options): ''' if 'pirvateKey' in options: - options['privateKey'] = self._upload(options['privateKey'])['filename'] + options['privateKey'] = ( + self._upload(options['privateKey'])['filename']) if 'publicKey' in options: - options['publicKey'] = self._upload(options['publicKey'])['filename'] + options['publicKey'] = ( + self._upload(options['publicKey'])['filename']) return self.raw_query("credential", "add", data=options) @@ -567,7 +590,8 @@ def credential_delete(self, *ids): }) def plugins(self, plugin_type='all', sort='id', direction='asc', - size=1000, offset=0, all=True, loops=0, since=None, **filterset): + size=1000, offset=0, all=True, loops=0, since=None, + **filterset): """plugins Returns a list of of the plugins and their associated families. For simplicity purposes, the plugin family names will be injected into the @@ -614,14 +638,14 @@ def plugins(self, plugin_type='all', sort='id', direction='asc', # Instance up and running to test... # --- # Next we convert the family dictionary list into a flat dictionary. - #fams = {} - #for famitem in data['families']: + # fams = {} + # for famitem in data['families']: # fams[famitem['id']] = famitem['name'] # Then we parse thtrough the data set, adding in the family name # into the plugin definition before adding it into the plugins list. for plugin in data['plugins']: - # plugin['familyName'] = fams[plugin['familyID']] + # plugin['familyName'] = fams[plugin['familyID']] plugins.append(plugin) # --- @@ -776,17 +800,18 @@ def scan_download(self, scan_id, format='v2'): 'downloadType': format, 'scanResultID': scan_id, } - data = self.raw_query('scanResult', 'download', data=payload, dejson=False) + data = self.raw_query('scanResult', 'download', data=payload, + dejson=False) bobj = StringIO() bobj.write(data) zfile = ZipFile(bobj) return zfile.read(zfile.namelist()[0]) - ### WARNING ### + # ### WARNING ### # All of the functions below are not part of the API documentation. This # means that it is entirely possible for one or all of these to change # without notification as they are not part of the documented API. - ############### + # ############## def _upload(self, fileobj): """_upload filename @@ -795,8 +820,9 @@ def _upload(self, fileobj): UN-DOCUMENTED CALL: This function is not considered stable. """ - return self.raw_query('file', 'upload', - data={'returnContent': 'false'}, files={'Filedata': fileobj}) + return self.raw_query('file', 'upload', + data={'returnContent': 'false'}, + files={'Filedata': fileobj}) def dashboard_import(self, name, fileobj): """dashboard_import Dashboard_Name, filename @@ -822,7 +848,6 @@ def report_import(self, name, filename): 'name': name, }) - def download_repository(self, repo_id): '''download_repository Repository_Id Download the tarball of the repository id specified. @@ -833,7 +858,6 @@ def download_repository(self, repo_id): 'id': repo_id }, dejson=False) - def asset_create(self, name, items, tag='', description='', atype='static'): '''asset_create_static name, ips, tags, description Create a new asset list with the defined information. @@ -862,7 +886,6 @@ def asset_create(self, name, items, tag='', description='', atype='static'): data['definedDNSNames'] = ' '.join(items) return self.raw_query('asset', 'add', data=data) - def asset_create_combo(self, name, combo, tag='', description=''): '''asset_create_combo name, combination, tag, description Creates a new combination asset list. Operands can be either asset list @@ -903,7 +926,6 @@ def asset_create_combo(self, name, combo, tag='', description=''): 'combinations': combo, }) - def risk_rule(self, rule_type, rule_value, port, proto, plugin_id, repo_ids, comment='', expires='-1', severity=None): '''accept_risk rule_type, rule_value, port, proto, plugin_id, comment @@ -943,11 +965,11 @@ def risk_rule(self, rule_type, rule_value, port, proto, plugin_id, data['expires'] = expires return self.raw_query('acceptRiskRule', 'add', data=data) else: - sevlevels = {'info': 0, 'low': 1, 'medium': 2, 'high': 3, 'critical': 4} + sevlevels = {'info': 0, 'low': 1, 'medium': 2, 'high': 3, + 'critical': 4} data['severity'] = sevlevels[severity] return self.raw_query('recastRiskRule', 'add', data=data) - def group_add(self, name, restrict, repos, lces=[], assets=[], queries=[], policies=[], dashboards=[], credentials=[], description=''): '''group_add name, restrict, repos diff --git a/securitycenter/sc5.py b/securitycenter/sc5.py index 206c60a..a85b8dc 100644 --- a/securitycenter/sc5.py +++ b/securitycenter/sc5.py @@ -1,9 +1,11 @@ -from .base import BaseAPI, APIError, logging +from .base import BaseAPI, APIError class SecurityCenter5(BaseAPI): _pre = 'rest/' - def __init__(self, host, port=443, ssl_verify=False, scheme='https', log=False, timeout=None): + + def __init__(self, host, port=443, ssl_verify=False, scheme='https', + log=False, timeout=None): '''SecurityCenter 5 API Wrapper This class is designed to handle authentication management for the SecurityCenter 5.x API. This is by no means a complete model of @@ -12,7 +14,8 @@ def __init__(self, host, port=443, ssl_verify=False, scheme='https', log=False, passes and there is a desire to develop them. For more information, please See Tenable's official API documentation - at: https://support.tenable.com/support-center/cerberus-support-center/includes/widgets/sc_api/index.html + at: https://support.tenable.com/support-center/cerberus-support-center/ + includes/widgets/sc_api/index.html ''' BaseAPI.__init__(self, host, port, ssl_verify, scheme, log, timeout) d = self.get('system').json() @@ -40,13 +43,14 @@ def _builder(self, **kwargs): return kwargs def login(self, user, passwd): - '''Logs the user into SecurityCenter and stores the needed token and cookies.''' + '''Logs the user into SecurityCenter and stores the needed token and + cookies.''' resp = self.post('token', json={'username': user, 'password': passwd}) self._token = resp.json()['response']['token'] def logout(self): '''Logs out of SecurityCenter and removed the cookies and token.''' - resp = self.delete('token') + self.delete('token') self._reset_session() def upload(self, fileobj): @@ -55,37 +59,52 @@ def upload(self, fileobj): def analysis(self, *filters, **kwargs): '''Analysis - A thin wrapper to handle vuln/event/mobile/log analysis through the API. This - function handles expanding multiple filters and will translate arbitrary arguments - into the format that SecurityCenter's analysis call expect them to be in. + A thin wrapper to handle vuln/event/mobile/log analysis through the API. + This function handles expanding multiple filters and will translate + arbitrary arguments into the format that SecurityCenter's analysis call + expect them to be in. - In order to make the filter expander more useful for SecurityCenter5 verses the - SecurityCenter4 class, filters are no longer done as kwargs, and instead are done - as a list of tuples. For example, to get the IP Summary of all of the hosts in - the 10.10.0.0/16 network, we would make the following call: + In order to make the filter expander more useful for SecurityCenter5 + verses the SecurityCenter4 class, filters are no longer done as kwargs, + and instead are done as a list of tuples. For example, to get the IP + Summary of all of the hosts in the 10.10.0.0/16 network, we would make + the following call: vulns = sc.analysis(('ip','=','10.10.0.0/16'), tool='sumip') - If multiple filters are desired, then it's simply a matter of entering multiple tuples. - - The Analysis function also has a few functions that are sligly deviated from the API - guides. All of these options are optional, however can significantly change the how - the API is being called. - - page - Default "all" The current page in the pagination sequence. - page_size - Default 1000 The page size (number of returned results) - page_obj - Default "return_results" The class thats called after every API pull. - page_kwargs - Default {} Any additional arguments that need to be passed - to the page_obj class when instantiated. - type - Default "vuln" This is the type of data that will be returned. - As all of the LCE and Vuln calls have been - collapsed into "analysis" in the API, each filter - needs to have these specified. This module does - that legwork for you. - sourceType - Default "cumulative" This is how we specify individual scans, LCE archive - silos, and other datasets that stand outside the norm. + If multiple filters are desired, then it's simply a matter of entering + multiple tuples. + + The Analysis function also has a few functions that are slightly + deviated from the API guides. All of these options are optional, + however can significantly change the how the API is being called. + + page - Default "all" + The current page in the pagination sequence. + + page_size - Default 1000 + The page size (number of returned results) + + page_obj - Default "return_results" + The class thats called after every API pull. + + page_kwargs - Default {} + Any additional arguments that need to be passed + to the page_obj class when instantiated. + + type - Default "vuln" + This is the type of data that will be returned. + As all of the LCE and Vuln calls have been + collapsed into "analysis" in the API, each filter + needs to have these specified. This module does + that legwork for you. + + sourceType - Default "cumulative" + This is how we specify individual scans, LCE archive + silos, and other datasets that stand outside the norm. ''' output = [] + def return_results(**kwargs): return kwargs['resp'].json()['response']['results'] @@ -93,41 +112,56 @@ def return_generator(**kwargs): for item in kwargs['resp'].json()['response']['results']: yield item - # These values are commonly used and/or are generally not changed from the default. - # If we do not see them specified by the user then we will need to add these in - # for later parsing... - if 'page' not in kwargs: kwargs['page'] = 'all' - if 'page_size' not in kwargs: kwargs['page_size'] = 1000 - if 'page_obj' not in kwargs: kwargs['page_obj'] = return_results - if 'page_kwargs' not in kwargs: kwargs['page_kwargs'] = {} - if 'type' not in kwargs: kwargs['type'] = 'vuln' - if 'sourceType' not in kwargs: kwargs['sourceType'] = 'cumulative' - if 'generator' in kwargs: kwargs['generator'] = return_generator - - # New we need to pull out the options from kwargs as we will be using hwargs as - # the basis for the query that will be sent to SecurityCenter. + # These values are commonly used and/or are generally not changed from + # the default. If we do not see them specified by the user then we will + # need to add these in for later parsing... + + if 'page' not in kwargs: + kwargs['page'] = 'all' + if 'page_size' not in kwargs: + kwargs['page_size'] = 1000 + if 'page_obj' not in kwargs: + kwargs['page_obj'] = return_results + if 'page_kwargs' not in kwargs: + kwargs['page_kwargs'] = {} + if 'type' not in kwargs: + kwargs['type'] = 'vuln' + if 'sourceType' not in kwargs: + kwargs['sourceType'] = 'cumulative' + if 'generator' in kwargs: + kwargs['generator'] = return_generator + + # New we need to pull out the options from kwargs as we will be using + # hwargs as the basis for the query that will be sent to SecurityCenter. + opts = {} - for opt in ['page', 'page_size', 'page_obj', 'page_kwargs', 'generator']: + for opt in ['page', 'page_size', 'page_obj', 'page_kwargs', + 'generator']: if opt in kwargs: opts[opt] = kwargs[opt] del kwargs[opt] - # If a query option was not set, then we will have to. The hope here is that - # we can make a lot of the pain of building the query pretty easy by simply - # accepting tuples of the filters. + # If a query option was not set, then we will have to. The hope here is + # that we can make a lot of the pain of building the query pretty easy + # by simply accepting tuples of the filters. + if 'query' not in kwargs: kwargs['query'] = { 'tool': kwargs['tool'], 'type': kwargs['type'], - 'filters': [{'filterName': f[0], 'operator': f[1], 'value': f[2], 'type': kwargs['type']} for f in filters] + 'filters': [{'filterName': f[0], 'operator': f[1], + 'value': f[2], + 'type': kwargs['type']} for f in filters] } del kwargs['tool'] if opts['page'] == 'all': kwargs['query']['startOffset'] = 0 kwargs['query']['endOffset'] = opts['page_size'] else: - kwargs['query']['startOffset'] = opts['page'] * opts['page_size'] - kwargs['query']['endOffset'] = (opts['page'] + 1) * opts['page_size'] + kwargs['query']['startOffset'] = (opts['page'] * + opts['page_size']) + kwargs['query']['endOffset'] = ((opts['page'] + 1) * + opts['page_size']) count = 0 total_records = opts['page_size']