def test_get_keys_by_pattern_error(self, mock_redis):
ctx = mock_redis.return_value
with self.assertRaises(RequiredArgument):
OpenvasDB.get_keys_by_pattern(None, 'a')
with self.assertRaises(RequiredArgument):
OpenvasDB.get_keys_by_pattern(ctx, None)
def test_get_list_item_error(self, mock_redis):
ctx = mock_redis.return_value
with self.assertRaises(RequiredArgument):
OpenvasDB.get_list_item(None, 'foo')
with self.assertRaises(RequiredArgument):
OpenvasDB.get_list_item(ctx, None)
def add_vt_to_cache(self, vt_id: str, vt: List[str]):
if not vt_id:
raise RequiredArgument('add_vt_to_cache', 'vt_id')
if not vt:
raise RequiredArgument('add_vt_to_cache', 'vt')
if not isinstance(vt, list) or len(vt) != 15:
raise OspdOpenvasError('Error trying to load the VT'
' {} in cache'.format(vt))
OpenvasDB.add_single_list(self.ctx, vt_id, vt)
def test_add_single_list(self, mock_redis):
ctx = mock_redis.return_value
pipeline = ctx.pipeline.return_value
pipeline.delete.return_value = None
pipeline.execute.return_value = (None, 0)
OpenvasDB.add_single_list(ctx, 'a', ['12', '11', '12'])
pipeline.delete.assert_called_once_with('a')
pipeline.rpush.assert_called_once_with('a', '12', '11', '12')
assert_called(pipeline.execute)
def test_create_context_fail(self, mock_redis):
mock_redis.side_effect = RCE
logging.Logger.error = MagicMock()
with patch.object(time, 'sleep', return_value=None):
with self.assertRaises(SystemExit):
OpenvasDB.create_context()
logging.Logger.error.assert_called_with( # pylint: disable=no-member
'Redis Error: Not possible to connect to the kb.')
def test_set_single_item(self, mock_redis):
ctx = mock_redis.return_value
pipeline = ctx.pipeline.return_value
pipeline.delete.return_value = None
pipeline.rpush.return_value = None
pipeline.execute.return_value = None
OpenvasDB.set_single_item(ctx, 'foo', ['bar'])
pipeline.delete.assert_called_once_with('foo')
pipeline.rpush.assert_called_once_with('foo', 'bar')
assert_called(pipeline.execute)
def add_vt_to_cache(self, vt_id: str, vt: List[str]):
if not vt_id:
raise RequiredArgument('add_vt_to_cache', 'vt_id')
if not vt:
raise RequiredArgument('add_vt_to_cache', 'vt')
if not isinstance(vt, list) or len(vt) != 15:
raise OspdOpenvasError(
f'Error trying to load the VT {vt} in cache')
OpenvasDB.add_single_list(self.ctx, vt_id, vt)
OpenvasDB.add_single_item(self.ctx, f'filename:{vt[0]}', [int(time())])
def get_nvt_refs(self, oid: str) -> Optional[Dict[str, str]]:
"""Get a full NVT.
Arguments:
oid: OID of VT from which to get the VT references.
Returns:
A dictionary with the VT references.
"""
resp = OpenvasDB.get_list_item(
self.ctx,
"nvt:%s" % oid,
start=NVT_META_FIELDS.index("NVT_CVES_POS"),
end=NVT_META_FIELDS.index("NVT_XREFS_POS"),
)
if not isinstance(resp, list) or len(resp) == 0:
return None
subelem = ['cve', 'bid', 'xref']
refs = dict()
for child, res in zip(subelem, resp):
refs[child] = res.split(", ")
return refs
def test_get_elem_pattern_by_index(self, mock_redis):
ctx = mock_redis.return_value
ctx.keys.return_value = ['aa', 'ab']
ctx.lindex.side_effect = [1, 2]
ret = OpenvasDB.get_elem_pattern_by_index(ctx, 'a')
self.assertEqual(list(ret), [('aa', 1), ('ab', 2)])
def test_get_last_list_item(self, mock_redis):
ctx = mock_redis.return_value
ctx.rpop.return_value = 'foo'
ret = OpenvasDB.get_last_list_item(ctx, 'name')
self.assertEqual(ret, 'foo')
ctx.rpop.assert_called_with('name')
def test_get_list_item(self, mock_redis):
ctx = mock_redis.return_value
ctx.lrange.return_value = ['1234']
ret = OpenvasDB.get_list_item(ctx, 'name')
self.assertEqual(ret, ['1234'])
assert_called(ctx.lrange)
def test_get_keys_by_pattern(self, mock_redis):
ctx = mock_redis.return_value
ctx.keys.return_value = ['nvt:2', 'nvt:1']
ret = OpenvasDB.get_keys_by_pattern(ctx, 'nvt:*')
# Return sorted list
self.assertEqual(ret, ['nvt:1', 'nvt:2'])
def test_get_single_item(self, mock_redis):
ctx = mock_redis.return_value
ctx.lindex.return_value = 'a'
value = OpenvasDB.get_single_item(ctx, 'a')
self.assertEqual(value, 'a')
ctx.lindex.assert_called_once_with('a', 0)
def test_get_filenames_and_oids(self, mock_redis):
ctx = mock_redis.return_value
ctx.keys.return_value = ['nvt:1', 'nvt:2']
ctx.lindex.side_effect = ['aa', 'ab']
ret = OpenvasDB.get_filenames_and_oids(ctx)
self.assertEqual(list(ret), [('aa', '1'), ('ab', '2')])
def test_find_database_by_pattern_none(self, mock_redis):
ctx = mock_redis.return_value
ctx.keys.return_value = None
new_ctx, index = OpenvasDB.find_database_by_pattern('foo*', 123)
self.assertIsNone(new_ctx)
self.assertIsNone(index)
def get_oids(self) -> Iterator[Tuple[str, str]]:
"""Get the list of NVT file names and OIDs.
Returns:
A i. Each single list contains the filename
as first element and the oid as second one.
"""
return OpenvasDB.get_filenames_and_oids(self.ctx)
def test_get_pattern(self, mock_redis):
ctx = mock_redis.return_value
ctx.keys.return_value = ['a', 'b']
ctx.lrange.return_value = [1, 2, 3]
ret = OpenvasDB.get_pattern(ctx, 'a')
self.assertEqual(ret, [['a', [1, 2, 3]], ['b', [1, 2, 3]]])
def test_get_key_count_with_default_pattern(self, mock_redis):
ctx = mock_redis.return_value
ctx.keys.return_value = ['aa', 'ab']
ret = OpenvasDB.get_key_count(ctx)
self.assertEqual(ret, 2)
ctx.keys.assert_called_with('*')
def test_get_db_connection(self, mock_openvas: MagicMock,
mock_redis: MagicMock):
OpenvasDB._db_address = None # pylint: disable=protected-access
mock_settings = mock_openvas.get_settings.return_value
mock_settings.get.return_value = None
self.assertIsNone(OpenvasDB.get_database_address())
# set the first time
mock_openvas.get_settings.return_value = {'db_address': '/foo/bar'}
self.assertEqual(OpenvasDB.get_database_address(), "/foo/bar")
self.assertEqual(mock_openvas.get_settings.call_count, 2)
# should cache address
self.assertEqual(OpenvasDB.get_database_address(), "/foo/bar")
self.assertEqual(mock_openvas.get_settings.call_count, 2)
def test_get_key_count(self, mock_redis):
ctx = mock_redis.return_value
ctx.keys.return_value = ['aa', 'ab']
ret = OpenvasDB.get_key_count(ctx, "foo")
self.assertEqual(ret, 2)
ctx.keys.assert_called_with('foo')
def test_find_database_by_pattern(self, mock_redis):
ctx = mock_redis.return_value
# keys is called twice per iteration
ctx.keys.side_effect = [None, None, None, None, True, True]
new_ctx, index = OpenvasDB.find_database_by_pattern('foo*', 123)
self.assertEqual(new_ctx, ctx)
self.assertEqual(index, 2)
def get_feed_version(self) -> Optional[str]:
"""Get feed version of the nvti cache db.
Returns the feed version or None if the nvt feed isn't available.
"""
if not self.ctx:
# no nvti cache db available yet
return None
return OpenvasDB.get_single_item(self.ctx, NVTI_CACHE_NAME)
def get_feed_version(self) -> Optional[str]:
"""Get feed version of the nvti cache db.
Returns the feed version or None if the nvt feed isn't available.
"""
if not self.ctx:
# no nvti cache db available yet
return None
# no feed version for notus otherwise tha would be a contract change
return OpenvasDB.get_single_item(self.ctx, NVTI_CACHE_NAME)
def get_nvt_prefs(self, oid: str) -> Optional[List[str]]:
"""Get NVT preferences.
Arguments:
ctx: Redis context to be used.
oid: OID of VT from which to get the VT preferences.
Returns:
A list with the VT preferences.
"""
key = 'oid:%s:prefs' % oid
return OpenvasDB.get_list_item(self.ctx, key)
def get_nvt_prefs(self, oid: str) -> Optional[List[str]]:
"""Get NVT preferences.
Arguments:
ctx: Redis context to be used.
oid: OID of VT from which to get the VT preferences.
Returns:
A list with the VT preferences.
"""
key = f'oid:{oid}:prefs'
# notus doesn't seem to have preferences, ignoring
return OpenvasDB.get_list_item(self.ctx, key)
def get_oids(self) -> Iterator[Tuple[str, str]]:
"""Get the list of NVT file names and OIDs.
Returns:
An iterable of tuples of file name and oid.
"""
if self.notus:
for f, oid in self.notus.get_filenames_and_oids():
yield (f, oid)
if self.ctx:
for f, oid in OpenvasDB.get_filenames_and_oids(self.ctx):
if not self.notus or not self.notus.exists(oid):
yield (f, oid)
def get_nvt_family(self, oid: str) -> str:
"""Get NVT family
Arguments:
oid: OID of VT from which to get the VT family.
Returns:
A str with the VT family.
"""
return OpenvasDB.get_single_item(
self.ctx,
'nvt:%s' % oid,
index=NVT_META_FIELDS.index("NVT_FAMILY_POS"),
)
def get_nvt_metadata(self, oid: str) -> Optional[Dict[str, str]]:
"""Get a full NVT. Returns an XML tree with the NVT metadata.
Arguments:
oid: OID of VT from which to get the metadata.
Returns:
A dictionary with the VT metadata.
"""
resp = OpenvasDB.get_list_item(
self.ctx,
f"nvt:{oid}",
start=NVT_META_FIELDS.index("NVT_FILENAME_POS"),
end=NVT_META_FIELDS.index("NVT_NAME_POS"),
)
if not isinstance(resp, list) or len(resp) == 0:
return None
subelem = [
'filename',
'required_keys',
'mandatory_keys',
'excluded_keys',
'required_udp_ports',
'required_ports',
'dependencies',
'tag',
'cve',
'bid',
'xref',
'category',
'family',
'name',
]
custom = dict()
custom['refs'] = dict()
for child, res in zip(subelem, resp):
if child not in ['cve', 'bid', 'xref', 'tag'] and res:
custom[child] = res
elif child == 'tag':
custom.update(self._parse_metadata_tags(res, oid))
elif child in ['cve', 'bid', 'xref'] and res:
custom['refs'][child] = res.split(", ")
custom['vt_params'] = dict()
custom['vt_params'].update(self.get_nvt_params(oid))
return custom
def test_pop_list_items_no_results(self, mock_redis):
ctx = mock_redis.return_value
pipeline = ctx.pipeline.return_value
pipeline.lrange.return_value = None
pipeline.delete.return_value = None
pipeline.execute.return_value = (None, 0)
ret = OpenvasDB.pop_list_items(ctx, 'foo')
self.assertEqual(ret, [])
pipeline.lrange.assert_called_once_with('foo', 0, -1)
pipeline.delete.assert_called_once_with('foo')
assert_called(pipeline.execute)
def test_pop_list_items_with_results(self, mock_redis):
ctx = mock_redis.return_value
pipeline = ctx.pipeline.return_value
pipeline.lrange.return_value = None
pipeline.delete.return_value = None
pipeline.execute.return_value = [['c', 'b', 'a'], 2]
ret = OpenvasDB.pop_list_items(ctx, 'results')
# reversed list
self.assertEqual(ret, ['a', 'b', 'c'])
pipeline.lrange.assert_called_once_with('results', 0, -1)
pipeline.delete.assert_called_once_with('results')
assert_called(pipeline.execute)