Use a session and don't overwrite imports

This commit is contained in:
pemontto
2019-04-08 12:24:22 +10:00
parent 5edde8760a
commit 33f2a5a3d1

View File

@ -34,7 +34,10 @@ class NessusAPI(object):
self.base = 'https://{hostname}:{port}'.format(hostname=hostname, port=port)
self.verbose = verbose
self.headers = {
self.session = requests.Session()
self.session.verify = False
self.session.stream = True
self.session.headers = {
'Origin': self.base,
'Accept-Encoding': 'gzip, deflate, br',
'Accept-Language': 'en-US,en;q=0.8',
@ -52,27 +55,24 @@ class NessusAPI(object):
self.scan_ids = self.get_scan_ids()
def login(self):
resp = self.get_token()
auth = '{"username":"%s", "password":"%s"}' % (self.user, self.password)
resp = self.request(self.SESSION, data=auth, json_output=False)
if resp.status_code == 200:
self.headers['X-Cookie'] = 'token={token}'.format(token=resp.json()['token'])
self.session.headers['X-Cookie'] = 'token={token}'.format(token=resp.json()['token'])
else:
raise Exception('[FAIL] Could not login to Nessus')
def request(self, url, data=None, headers=None, method='POST', download=False, json=False):
if headers is None:
headers = self.headers
def request(self, url, data=None, headers=None, method='POST', download=False, json_output=False):
timeout = 0
success = False
method = method.lower()
url = self.base + url
self.logger.debug('Requesting to url {}'.format(url))
methods = {'GET': requests.get,
'POST': requests.post,
'DELETE': requests.delete}
while (timeout <= 10) and (not success):
data = methods[method](url, data=data, headers=self.headers, verify=False)
if data.status_code == 401:
response = getattr(self.session, method)(url, data=data)
if response.status_code == 401:
if url == self.base + self.SESSION:
break
try:
@ -84,20 +84,15 @@ class NessusAPI(object):
else:
success = True
if json:
data = data.json()
if json_output:
return response.json()
if download:
self.logger.debug('Returning data.content')
return data.content
return data
def get_token(self):
auth = '{"username":"%s", "password":"%s"}' % (self.user, self.password)
token = self.request(self.SESSION, data=auth, json=False)
return token
return response.content
return response
def get_scans(self):
scans = self.request(self.SCANS, method='GET', json=True)
scans = self.request(self.SCANS, method='GET', json_output=True)
return scans
def get_scan_ids(self):
@ -107,10 +102,10 @@ class NessusAPI(object):
return scan_ids
def get_scan_history(self, scan_id):
data = self.request(self.SCAN_ID.format(scan_id=scan_id), method='GET', json=True)
data = self.request(self.SCAN_ID.format(scan_id=scan_id), method='GET', json_output=True)
return data['history']
def download_scan(self, scan_id=None, history=None, export_format="", chapters="", dbpasswd="", profile=""):
def download_scan(self, scan_id=None, history=None, export_format="", profile=""):
running = True
counter = 0
@ -120,7 +115,7 @@ class NessusAPI(object):
else:
query = self.EXPORT_HISTORY.format(scan_id=scan_id, history_id=history)
scan_id = str(scan_id)
req = self.request(query, data=json.dumps(data), method='POST', json=True)
req = self.request(query, data=json.dumps(data), method='POST', json_output=True)
try:
file_id = req['file']
token_id = req['token'] if 'token' in req else req['temp_token']
@ -131,7 +126,7 @@ class NessusAPI(object):
time.sleep(2)
counter += 2
report_status = self.request(self.EXPORT_STATUS.format(scan_id=scan_id, file_id=file_id), method='GET',
json=True)
json_output=True)
running = report_status['status'] != 'ready'
sys.stdout.write(".")
sys.stdout.flush()