Simplify pandas operations, update transforms
This commit is contained in:
@ -203,15 +203,11 @@ class NessusAPI(object):
|
||||
df.drop('CVSS', axis=1, inplace=True)
|
||||
df.drop('IP Address', axis=1, inplace=True)
|
||||
|
||||
# Map fields from COLUMN_MAPPING
|
||||
fields = [x.lower() for x in df.columns]
|
||||
for field, replacement in self.COLUMN_MAPPING.iteritems():
|
||||
if field in fields:
|
||||
self.logger.debug('Renaming "{}" to "{}"'.format(field, replacement))
|
||||
fields[fields.index(field)] = replacement
|
||||
# Lowercase and map fields from COLUMN_MAPPING
|
||||
df.columns = [x.lower() for x in df.columns]
|
||||
df.rename(columns=self.COLUMN_MAPPING, inplace=True)
|
||||
df.columns = [x.replace(' ', '_') for x in df.columns]
|
||||
|
||||
fields = [x.replace(' ', '_') for x in fields]
|
||||
df.columns = fields
|
||||
return df
|
||||
|
||||
def transform_values(self, df):
|
||||
@ -227,8 +223,7 @@ class NessusAPI(object):
|
||||
|
||||
# Map risk to a SEVERITY MAPPING value
|
||||
self.logger.debug('Mapping risk to severity number')
|
||||
df['risk_number'] = df['risk'].str.lower()
|
||||
df['risk_number'].replace(self.SEVERITY_MAPPING, inplace=True)
|
||||
df['risk_number'] = df['risk'].str.lower().map(self.SEVERITY_MAPPING)
|
||||
|
||||
if self.profile == 'tenable':
|
||||
self.logger.debug('Combinging CVSS vectors for tenable')
|
||||
@ -243,9 +238,14 @@ class NessusAPI(object):
|
||||
.apply(lambda x: '{}/{}'.format(x[0], x[1]), axis=1)
|
||||
.str.rstrip('/nan')
|
||||
)
|
||||
# CVSS score = cvss_temporal if cvss_temporal else cvss_base
|
||||
|
||||
df.drop(['cvss_temporal_vector', 'cvss3_temporal_vector'], axis=1, inplace=True)
|
||||
|
||||
# CVSS score = cvss3_temporal or cvss3_base or cvss_temporal or cvss_base
|
||||
df['cvss'] = df['cvss_base']
|
||||
df.loc[df['cvss_temporal'].notnull(), 'cvss'] = df['cvss_temporal']
|
||||
df['cvss3'] = df['cvss3_base']
|
||||
df.loc[df['cvss3_temporal'].notnull(), 'cvss3'] = df['cvss3_temporal']
|
||||
|
||||
df.fillna('', inplace=True)
|
||||
return df
|
@ -144,15 +144,10 @@ class qualysVulnScan:
|
||||
def map_fields(self, df):
|
||||
self.logger.info('Mapping fields')
|
||||
|
||||
# Map fields from COLUMN_MAPPING
|
||||
fields = [x.lower() for x in df.columns]
|
||||
for field, replacement in self.COLUMN_MAPPING.iteritems():
|
||||
if field in fields:
|
||||
self.logger.info('Renaming "{}" to "{}"'.format(field, replacement))
|
||||
fields[fields.index(field)] = replacement
|
||||
|
||||
fields = [x.replace(' ', '_') for x in fields]
|
||||
df.columns = fields
|
||||
# Lowercase and map fields from COLUMN_MAPPING
|
||||
df.columns = [x.lower() for x in df.columns]
|
||||
df.rename(columns=self.COLUMN_MAPPING, inplace=True)
|
||||
df.columns = [x.replace(' ', '_') for x in df.columns]
|
||||
|
||||
return df
|
||||
|
||||
@ -165,32 +160,28 @@ class qualysVulnScan:
|
||||
df['protocol'] = df['protocol'].str.lower()
|
||||
|
||||
# Contruct the CVSS vector
|
||||
df['cvss_vector'] = ''
|
||||
df.loc[df['cvss_base'].notnull(), 'cvss_vector'] = (
|
||||
df['cvss_vector'] = (
|
||||
df.loc[df['cvss_base'].notnull(), 'cvss_base']
|
||||
.str.split()
|
||||
.apply(lambda x: x[1])
|
||||
.str.replace('(', '')
|
||||
.str.replace(')', '')
|
||||
.str.strip('()')
|
||||
)
|
||||
df.loc[df['cvss_base'].notnull(), 'cvss_base'] = (
|
||||
df['cvss_base'] = (
|
||||
df.loc[df['cvss_base'].notnull(), 'cvss_base']
|
||||
.str.split()
|
||||
.apply(lambda x: x[0])
|
||||
)
|
||||
df['cvss_temporal_vector'] = ''
|
||||
df.loc[df['cvss_temporal'].notnull(), 'cvss_temporal_vector'] = (
|
||||
|
||||
df['cvss_temporal_vector'] = (
|
||||
df.loc[df['cvss_temporal'].notnull(), 'cvss_temporal']
|
||||
.str.split()
|
||||
.apply(lambda x: x[1])
|
||||
.str.replace('(', '')
|
||||
.str.replace(')', '')
|
||||
.str.strip('()')
|
||||
)
|
||||
df.loc[df['cvss_temporal'].notnull(), 'cvss_temporal'] = (
|
||||
df['cvss_temporal'] = (
|
||||
df.loc[df['cvss_temporal'].notnull(), 'cvss_temporal']
|
||||
.str.split()
|
||||
.apply(lambda x: x[0])
|
||||
.fillna('')
|
||||
)
|
||||
|
||||
# Combine base and temporal
|
||||
@ -198,7 +189,6 @@ class qualysVulnScan:
|
||||
df[['cvss_vector', 'cvss_temporal_vector']]
|
||||
.apply(lambda x: '{}/{}'.format(x[0], x[1]), axis=1)
|
||||
.str.rstrip('/nan')
|
||||
.fillna('')
|
||||
)
|
||||
|
||||
df.drop('cvss_temporal_vector', axis=1, inplace=True)
|
||||
@ -206,8 +196,11 @@ class qualysVulnScan:
|
||||
# Convert Qualys severity to standardised risk number
|
||||
df['risk_number'] = df['severity'].astype(int)-1
|
||||
|
||||
# CVSS score = cvss3_temporal or cvss3_base or cvss_temporal or cvss_base
|
||||
df['cvss'] = df['cvss_base']
|
||||
df.loc[df['cvss_temporal'].notnull(), 'cvss'] = df['cvss_temporal']
|
||||
df['cvss3'] = df['cvss3_base']
|
||||
df.loc[df['cvss3_temporal'].notnull(), 'cvss3'] = df['cvss3_temporal']
|
||||
|
||||
df.fillna('', inplace=True)
|
||||
|
||||
|
@ -84,7 +84,7 @@ class vulnWhispererBase(object):
|
||||
self.cur = self.conn.cursor()
|
||||
self.logger.info('Connected to database at {loc}'.format(loc=self.database))
|
||||
except Exception as e:
|
||||
self.logger.error('Could not connect to database at {loc}\nReason: {e} - Please ensure the path exist'.format(
|
||||
self.logger.error('Could not connect to database at {loc}\nReason: {e} - Please ensure the path exists'.format(
|
||||
e=e,
|
||||
loc=self.database))
|
||||
else:
|
||||
@ -189,7 +189,7 @@ class vulnWhispererBase(object):
|
||||
scan=self.write_path.encode('utf8')))
|
||||
else:
|
||||
os.path.exists(self.write_path)
|
||||
self.logger.info('Directory already exist for {scan} - Skipping creation'.format(
|
||||
self.logger.info('Directory already exists for {scan} - Skipping creation'.format(
|
||||
scan=self.write_path.encode('utf8')))
|
||||
|
||||
def get_latest_results(self, source, scan_name):
|
||||
@ -376,7 +376,7 @@ class vulnWhispererNessus(vulnWhispererBase):
|
||||
os.makedirs(self.path_check(f['name']))
|
||||
else:
|
||||
os.path.exists(self.path_check(f['name']))
|
||||
self.logger.info('Directory already exist for {scan} - Skipping creation'.format(
|
||||
self.logger.info('Directory already exists for {scan} - Skipping creation'.format(
|
||||
scan=self.path_check(f['name']).encode('utf8')))
|
||||
|
||||
# try download and save scans into each folder the belong to
|
||||
@ -419,7 +419,7 @@ class vulnWhispererNessus(vulnWhispererBase):
|
||||
|
||||
if os.path.isfile(relative_path_name):
|
||||
if self.develop:
|
||||
csv_in = pd.read_csv(relative_path_name)
|
||||
csv_in = pd.read_json(relative_path_name, lines=True)
|
||||
record_meta = (
|
||||
scan_name,
|
||||
scan_id,
|
||||
@ -433,7 +433,7 @@ class vulnWhispererNessus(vulnWhispererBase):
|
||||
0,
|
||||
)
|
||||
self.record_insert(record_meta)
|
||||
self.logger.info('File {filename} already exist! Updating database'.format(filename=relative_path_name))
|
||||
self.logger.info('File {filename} already exists! Updating database'.format(filename=relative_path_name))
|
||||
else:
|
||||
try:
|
||||
file_req = \
|
||||
@ -608,7 +608,7 @@ class vulnWhispererQualys(vulnWhispererBase):
|
||||
0,
|
||||
)
|
||||
self.record_insert(record_meta)
|
||||
self.logger.info('File {filename} already exist! Updating database'.format(filename=relative_path_name))
|
||||
self.logger.info('File {filename} already exists! Updating database'.format(filename=relative_path_name))
|
||||
|
||||
else:
|
||||
self.logger.info('Generating report for {}'.format(report_id))
|
||||
@ -775,7 +775,7 @@ class vulnWhispererOpenVAS(vulnWhispererBase):
|
||||
0,
|
||||
)
|
||||
self.record_insert(record_meta)
|
||||
self.logger.info('File {filename} already exist! Updating database'.format(filename=relative_path_name))
|
||||
self.logger.info('File {filename} already exists! Updating database'.format(filename=relative_path_name))
|
||||
|
||||
record_meta = (
|
||||
scan_name,
|
||||
@ -889,7 +889,7 @@ class vulnWhispererQualysVuln(vulnWhispererBase):
|
||||
0,
|
||||
)
|
||||
self.record_insert(record_meta)
|
||||
self.logger.info('File {filename} already exist! Updating database'.format(filename=relative_path_name))
|
||||
self.logger.info('File {filename} already exists! Updating database'.format(filename=relative_path_name))
|
||||
|
||||
else:
|
||||
try:
|
||||
|
Reference in New Issue
Block a user