ospd_openvas.db.KbDB.add_scan_id() - Code Metrics - greenbone/ospd-openvas - Measure and Improve Code Quality continuously with Scrutinizer

ospd_openvas.db.KbDB.add_scan_id() A
last analyzed 2021-07-12 05:11 UTC

↳ Parent: ospd_openvas.db

Complexity

Conditions

Size

Total Lines	3
Code Lines	3

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
eloc	3
nop	2
dl	0
loc	3
rs	10
c	0
b	0
f	0

# -*- coding: utf-8 -*-
# Copyright (C) 2014-2021 Greenbone Networks GmbH
#
# SPDX-License-Identifier: AGPL-3.0-or-later
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.


""" Access management for redis-based OpenVAS Scanner Database."""
import logging
import sys
import time

from typing import List, NewType, Optional, Iterable, Iterator, Tuple

import redis

from ospd.errors import RequiredArgument
from ospd_openvas.errors import OspdOpenvasError
from ospd_openvas.openvas import Openvas

SOCKET_TIMEOUT = 60  # in seconds
LIST_FIRST_POS = 0
LIST_LAST_POS = -1
LIST_ALL = 0

# Possible positions of nvt values in cache list.
NVT_META_FIELDS = [
    "NVT_FILENAME_POS",
    "NVT_REQUIRED_KEYS_POS",
    "NVT_MANDATORY_KEYS_POS",
    "NVT_EXCLUDED_KEYS_POS",
    "NVT_REQUIRED_UDP_PORTS_POS",
    "NVT_REQUIRED_PORTS_POS",
    "NVT_DEPENDENCIES_POS",
    "NVT_TAGS_POS",
    "NVT_CVES_POS",
    "NVT_BIDS_POS",
    "NVT_XREFS_POS",
    "NVT_CATEGORY_POS",
    "NVT_TIMEOUT_POS",
    "NVT_FAMILY_POS",
    "NVT_NAME_POS",
]

# Name of the namespace usage bitmap in redis.
DBINDEX_NAME = "GVM.__GlobalDBIndex"

logger = logging.getLogger(__name__)

# Types
RedisCtx = NewType('RedisCtx', redis.Redis)


class OpenvasDB:
    """Class to connect to redis, to perform queries, and to move
    from a KB to another."""

    _db_address = None

    @classmethod
    def get_database_address(cls) -> Optional[str]:
        if not cls._db_address:
            settings = Openvas.get_settings()

            cls._db_address = settings.get('db_address')

        return cls._db_address

    @classmethod
    def create_context(
        cls, dbnum: Optional[int] = 0, encoding: Optional[str] = 'latin-1'
    ) -> RedisCtx:
        """Connect to redis to the given database or to the default db 0 .

        Arguments:
            dbnum: The db number to connect to.
            encoding: The encoding to be used to read and write.

        Return a new redis context on success.
        """
        tries = 5
        while tries:
            try:
                ctx = redis.Redis(
                    unix_socket_path=cls.get_database_address(),
                    db=dbnum,
                    socket_timeout=SOCKET_TIMEOUT,
                    encoding=encoding,
                    decode_responses=True,
                )
                ctx.keys("test")
            except (redis.exceptions.ConnectionError, FileNotFoundError) as err:
                logger.debug(
                    'Redis connection lost: %s. Trying again in 5 seconds.', err
                )
                tries = tries - 1
                time.sleep(5)
                continue
            break

        if not tries:
            logger.error('Redis Error: Not possible to connect to the kb.')
            sys.exit(1)

        return ctx


    @classmethod
    def find_database_by_pattern(
        cls, pattern: str, max_database_index: int
    ) -> Tuple[Optional[RedisCtx], Optional[int]]:
        """Search a pattern inside all kbs up to max_database_index.

        Returns the redis context for the db and its index as a tuple or
        None, None if the db with the pattern couldn't be found.
        """
        for i in range(0, max_database_index):
            ctx = cls.create_context(i)
            if ctx.keys(pattern):
                return (ctx, i)

        return (None, None)

    @staticmethod
    def select_database(ctx: RedisCtx, kbindex: str):
        """Use an existent redis connection and select a redis kb.

        Arguments:
            ctx: Redis context to use.
            kbindex: The new kb to select
        """
        if not ctx:
            raise RequiredArgument('select_database', 'ctx')
        if not kbindex:
            raise RequiredArgument('select_database', 'kbindex')

        ctx.execute_command('SELECT ' + str(kbindex))

    @staticmethod
    def get_list_item(
        ctx: RedisCtx,
        name: str,
        start: Optional[int] = LIST_FIRST_POS,
        end: Optional[int] = LIST_LAST_POS,
    ) -> Optional[list]:
        """Returns the specified elements from `start` to `end` of the
        list stored as `name`.

        Arguments:
            ctx: Redis context to use.
            name: key name of a list.
            start: first range element to get.
            end: last range element to get.

        Return List specified elements in the key.
        """
        if not ctx:
            raise RequiredArgument('get_list_item', 'ctx')
        if not name:
            raise RequiredArgument('get_list_item', 'name')

        return ctx.lrange(name, start, end)

    @staticmethod
    def get_last_list_item(ctx: RedisCtx, name: str) -> str:
        if not ctx:
            raise RequiredArgument('get_last_list_item', 'ctx')
        if not name:
            raise RequiredArgument('get_last_list_item', 'name')

        return ctx.rpop(name)

    @staticmethod
    def pop_list_items(ctx: RedisCtx, name: str) -> List[str]:
        if not ctx:
            raise RequiredArgument('pop_list_items', 'ctx')
        if not name:
            raise RequiredArgument('pop_list_items', 'name')

        pipe = ctx.pipeline()
        pipe.lrange(name, LIST_FIRST_POS, LIST_LAST_POS)
        pipe.delete(name)
        results, redis_return_code = pipe.execute()

        # The results are left-pushed. To preserver the order
        # the result list must be reversed.
        if redis_return_code:
            results.reverse()
        else:
            results = []

        return results

    @staticmethod
    def get_key_count(ctx: RedisCtx, pattern: Optional[str] = None) -> int:
        """Get the number of keys matching with the pattern.

        Arguments:
            ctx: Redis context to use.
            pattern: pattern used as filter.
        """
        if not pattern:
            pattern = "*"

        if not ctx:
            raise RequiredArgument('get_key_count', 'ctx')

        return len(ctx.keys(pattern))

    @staticmethod
    def remove_list_item(ctx: RedisCtx, key: str, value: str):
        """Remove item from the key list.

        Arguments:
            ctx: Redis context to use.
            key: key name of a list.
            value: Value to be removed from the key.
        """
        if not ctx:
            raise RequiredArgument('remove_list_item ', 'ctx')
        if not key:
            raise RequiredArgument('remove_list_item', 'key')
        if not value:
            raise RequiredArgument('remove_list_item ', 'value')

        ctx.lrem(key, count=LIST_ALL, value=value)

    @staticmethod
    def get_single_item(
        ctx: RedisCtx,
        name: str,
        index: Optional[int] = LIST_FIRST_POS,
    ) -> Optional[str]:
        """Get a single KB element.

        Arguments:
            ctx: Redis context to use.
            name: key name of a list.
            index: index of the element to be return.
                   Defaults to the first element in the list.

        Return the first element of the list or None if the name couldn't be
        found.
        """
        if not ctx:
            raise RequiredArgument('get_single_item', 'ctx')
        if not name:
            raise RequiredArgument('get_single_item', 'name')

        return ctx.lindex(name, index)

    @staticmethod
    def add_single_list(ctx: RedisCtx, name: str, values: Iterable):
        """Add a single KB element with one or more values.
        The values can be repeated. If the key already exists will
        be removed an completely replaced.

        Arguments:
            ctx: Redis context to use.
            name: key name of a list.
            value: Elements to add to the key.
        """
        if not ctx:
            raise RequiredArgument('add_single_list', 'ctx')
        if not name:
            raise RequiredArgument('add_single_list', 'name')
        if not values:
            raise RequiredArgument('add_single_list', 'value')

        pipe = ctx.pipeline()
        pipe.delete(name)
        pipe.rpush(name, *values)
        pipe.execute()

    @staticmethod
    def add_single_item(ctx: RedisCtx, name: str, values: Iterable):
        """Add a single KB element with one or more values. Don't add
        duplicated values during this operation, but if the the same
        values already exists under the key, this will not be overwritten.

        Arguments:
            ctx: Redis context to use.
            name: key name of a list.
            value: Elements to add to the key.
        """
        if not ctx:
            raise RequiredArgument('add_single_item', 'ctx')
        if not name:
            raise RequiredArgument('add_single_item', 'name')
        if not values:
            raise RequiredArgument('add_single_item', 'value')

        ctx.rpush(name, *set(values))

    @staticmethod
    def set_single_item(ctx: RedisCtx, name: str, value: Iterable):
        """Set (replace) a single KB element. If the same key exists
        in the kb, it is completed removed. Values added are unique.

        Arguments:
            ctx: Redis context to use.
            name: key name of a list.
            value: New elements to add to the key.
        """
        if not ctx:
            raise RequiredArgument('set_single_item', 'ctx')
        if not name:
            raise RequiredArgument('set_single_item', 'name')
        if not value:
            raise RequiredArgument('set_single_item', 'value')

        pipe = ctx.pipeline()
        pipe.delete(name)
        pipe.rpush(name, *set(value))
        pipe.execute()

    @staticmethod
    def get_pattern(ctx: RedisCtx, pattern: str) -> List:
        """Get all items stored under a given pattern.

        Arguments:
            ctx: Redis context to use.
            pattern: key pattern to match.

        Return a list with the elements under the matched key.
        """
        if not ctx:
            raise RequiredArgument('get_pattern', 'ctx')
        if not pattern:
            raise RequiredArgument('get_pattern', 'pattern')

        items = ctx.keys(pattern)

        elem_list = []
        for item in items:
            elem_list.append(
                [
                    item,
                    ctx.lrange(item, start=LIST_FIRST_POS, end=LIST_LAST_POS),
                ]
            )
        return elem_list

    @classmethod
    def get_keys_by_pattern(cls, ctx: RedisCtx, pattern: str) -> List[str]:
        """Get all items with index 'index', stored under
        a given pattern.

        Arguments:
            ctx: Redis context to use.
            pattern: key pattern to match.

        Return a sorted list with the elements under the matched key
        """
        if not ctx:
            raise RequiredArgument('get_elem_pattern_by_index', 'ctx')
        if not pattern:
            raise RequiredArgument('get_elem_pattern_by_index', 'pattern')

        return sorted(ctx.keys(pattern))

    @classmethod
    def get_filenames_and_oids(
        cls,
        ctx: RedisCtx,
    ) -> Iterable[Tuple[str, str]]:
        """Get all items with index 'index', stored under
        a given pattern.

        Arguments:
            ctx: Redis context to use.

        Return an iterable where each single tuple contains the filename
            as first element and the oid as the second one.
        """
        if not ctx:
            raise RequiredArgument('get_filenames_and_oids', 'ctx')

        items = cls.get_keys_by_pattern(ctx, 'nvt:*')

        return ((ctx.lindex(item, 0), item[4:]) for item in items)



class BaseDB:
    def __init__(self, kbindex: int, ctx: Optional[RedisCtx] = None):
        if ctx is None:
            self.ctx = OpenvasDB.create_context(kbindex)
        else:
            self.ctx = ctx

        self.index = kbindex

    def flush(self):
        """Flush the database"""
        self.ctx.flushdb()


class BaseKbDB(BaseDB):
    def _add_single_item(
        self, name: str, values: Iterable, utf8_enc: Optional[bool] = False
    ):
        """Changing the encoding format of an existing redis context
        is not possible. Therefore a new temporary redis context is
        created to store key-values encoded with utf-8."""
        if utf8_enc:
            ctx = OpenvasDB.create_context(self.index, encoding='utf-8')
            OpenvasDB.add_single_item(ctx, name, values)
        else:
            OpenvasDB.add_single_item(self.ctx, name, values)

    def _set_single_item(self, name: str, value: Iterable):
        """Set (replace) a single KB element.

        Arguments:
            name: key name of a list.
            value: New elements to add to the key.
        """
        OpenvasDB.set_single_item(self.ctx, name, value)

    def _get_single_item(self, name: str) -> Optional[str]:
        """Get a single KB element.

        Arguments:
            name: key name of a list.
        """
        return OpenvasDB.get_single_item(self.ctx, name)

    def _get_list_item(
        self,
        name: str,
    ) -> Optional[List]:
        """Returns the specified elements from `start` to `end` of the
        list stored as `name`.

        Arguments:
            name: key name of a list.

        Return List specified elements in the key.
        """
        return OpenvasDB.get_list_item(self.ctx, name)

    def _pop_list_items(self, name: str) -> List:
        return OpenvasDB.pop_list_items(self.ctx, name)

    def _remove_list_item(self, key: str, value: str):
        """Remove item from the key list.

        Arguments:
            key: key name of a list.
            value: Value to be removed from the key.
        """
        OpenvasDB.remove_list_item(self.ctx, key, value)

    def get_result(self) -> Optional[str]:
        """Get and remove the oldest result from the list.

        Return the oldest scan results
        """
        return self._pop_list_items("internal/results")

    def get_status(self, openvas_scan_id: str) -> Optional[str]:
        """Return the status of the host scan"""
        return self._get_single_item('internal/{}'.format(openvas_scan_id))

    def __repr__(self):
        return '<{} index={}>'.format(self.__class__.__name__, self.index)


class ScanDB(BaseKbDB):
    """Database for a scanning a single host"""

    def select(self, kbindex: int) -> "ScanDB":
        """Select a redis kb.

        Arguments:
            kbindex: The new kb to select
        """
        OpenvasDB.select_database(self.ctx, kbindex)
        self.index = kbindex
        return self


class KbDB(BaseKbDB):
    def get_scan_databases(self) -> Iterator[ScanDB]:
        """Returns an iterator yielding corresponding ScanDBs

        The returned Iterator can't be converted to an Iterable like a List.
        Each yielded ScanDB must be used independently in a for loop. If the
        Iterator gets converted into an Iterable all returned ScanDBs will use
        the same redis context pointing to the same redis database.
        """
        dbs = self._get_list_item('internal/dbindex')
        scan_db = ScanDB(self.index)
        for kbindex in dbs:
            if kbindex == self.index:
                continue

            yield scan_db.select(kbindex)

    def add_scan_id(self, scan_id: str):
        self._add_single_item('internal/{}'.format(scan_id), ['new'])
        self._add_single_item('internal/scanid', [scan_id])

    def add_scan_preferences(self, openvas_scan_id: str, preferences: Iterable):
        self._add_single_item(
            'internal/{}/scanprefs'.format(openvas_scan_id), preferences
        )

    def add_credentials_to_scan_preferences(
        self, openvas_scan_id: str, preferences: Iterable
    ):
        """Force the usage of the utf-8 encoding, since some credentials
        contain special chars not supported by latin-1 encoding."""
        self._add_single_item(
            'internal/{}/scanprefs'.format(openvas_scan_id),
            preferences,
            utf8_enc=True,
        )

    def add_scan_process_id(self, pid: int):
        self._add_single_item('internal/ovas_pid', [pid])

    def get_scan_process_id(self) -> Optional[str]:
        return self._get_single_item('internal/ovas_pid')

    def remove_scan_database(self, scan_db: ScanDB):
        self._remove_list_item('internal/dbindex', scan_db.index)

    def target_is_finished(self, scan_id: str) -> bool:
        """Check if a target has finished."""

        status = self._get_single_item('internal/{}'.format(scan_id))

        if status is None:
            logger.error(
                "%s: Target set as finished because redis returned None as "
                "scanner status.",
                scan_id,
            )

        return status == 'finished' or status is None

    def stop_scan(self, openvas_scan_id: str):
        self._set_single_item(
            'internal/{}'.format(openvas_scan_id), ['stop_all']
        )

    def scan_is_stopped(self, scan_id: str) -> bool:
        """Check if the scan should be stopped"""
        status = self._get_single_item('internal/%s' % scan_id)
        return status == 'stop_all'

    def get_scan_status(self) -> List:
        """Get and remove the oldest host scan status from the list.

        Return a string which represents the host scan status.
        """
        return self._pop_list_items("internal/status")


class MainDB(BaseDB):
    """Main Database"""

    DEFAULT_INDEX = 0

    def __init__(self, ctx=None):
        super().__init__(self.DEFAULT_INDEX, ctx)

        self._max_dbindex = None

    @property
    def max_database_index(self):
        """Set the number of databases have been configured into kbr struct."""
        if self._max_dbindex is None:
            resp = self.ctx.config_get('databases')

            if len(resp) == 1:
                self._max_dbindex = int(resp.get('databases'))
            else:
                raise OspdOpenvasError(
                    'Redis Error: Not possible to get max_dbindex.'
                ) from None

        return self._max_dbindex

    def try_database(self, index: int) -> bool:
        """Check if a redis db is already in use. If not, set it
        as in use and return.

        Arguments:
            ctx: Redis object connected to the kb with the
                DBINDEX_NAME key.
            index: Number intended to be used.

        Return True if it is possible to use the db. False if the given db
            number is already in use.
        """
        _in_use = 1
        try:
            resp = self.ctx.hsetnx(DBINDEX_NAME, index, _in_use)
        except:
            raise OspdOpenvasError(
                'Redis Error: Not possible to set %s.' % DBINDEX_NAME
            ) from None

        return resp == 1

    def get_new_kb_database(self) -> Optional[KbDB]:
        """Return a new kb db to an empty kb."""
        for index in range(1, self.max_database_index):
            if self.try_database(index):
                kbdb = KbDB(index)
                kbdb.flush()
                return kbdb

        return None

    def find_kb_database_by_scan_id(
        self, scan_id: str
    ) -> Tuple[Optional[str], Optional["KbDB"]]:
        """Find a kb db by via a scan id"""
        for index in range(1, self.max_database_index):
            ctx = OpenvasDB.create_context(index)
            if OpenvasDB.get_key_count(ctx, 'internal/{}'.format(scan_id)):
                return KbDB(index, ctx)

        return None

    def release_database(self, database: BaseDB):
        self.release_database_by_index(database.index)
        database.flush()

    def release_database_by_index(self, index: int):
        self.ctx.hdel(DBINDEX_NAME, index)

    def release(self):
        self.release_database(self)


1			# -- coding: utf-8 --
2			# Copyright (C) 2014-2021 Greenbone Networks GmbH
3			#
4			# SPDX-License-Identifier: AGPL-3.0-or-later
5			#
6			# This program is free software: you can redistribute it and/or modify
7			# it under the terms of the GNU Affero General Public License as
8			# published by the Free Software Foundation, either version 3 of the
9			# License, or (at your option) any later version.
10			#
11			# This program is distributed in the hope that it will be useful,
12			# but WITHOUT ANY WARRANTY; without even the implied warranty of
13			# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14			# GNU Affero General Public License for more details.
15			#
16			# You should have received a copy of the GNU Affero General Public License
17			# along with this program. If not, see <http://www.gnu.org/licenses/>.
18
19
20			""" Access management for redis-based OpenVAS Scanner Database."""
21			import logging
22			import sys
23			import time
24
25			from typing import List, NewType, Optional, Iterable, Iterator, Tuple
26
27			import redis
28
29			from ospd.errors import RequiredArgument
30			from ospd_openvas.errors import OspdOpenvasError
31			from ospd_openvas.openvas import Openvas
32
33			SOCKET_TIMEOUT = 60 # in seconds
34			LIST_FIRST_POS = 0
35			LIST_LAST_POS = -1
36			LIST_ALL = 0
37
38			# Possible positions of nvt values in cache list.
39			NVT_META_FIELDS = [
40			"NVT_FILENAME_POS",
41			"NVT_REQUIRED_KEYS_POS",
42			"NVT_MANDATORY_KEYS_POS",
43			"NVT_EXCLUDED_KEYS_POS",
44			"NVT_REQUIRED_UDP_PORTS_POS",
45			"NVT_REQUIRED_PORTS_POS",
46			"NVT_DEPENDENCIES_POS",
47			"NVT_TAGS_POS",
48			"NVT_CVES_POS",
49			"NVT_BIDS_POS",
50			"NVT_XREFS_POS",
51			"NVT_CATEGORY_POS",
52			"NVT_TIMEOUT_POS",
53			"NVT_FAMILY_POS",
54			"NVT_NAME_POS",
55			]
56
57			# Name of the namespace usage bitmap in redis.
58			DBINDEX_NAME = "GVM.__GlobalDBIndex"
59
60			logger = logging.getLogger(__name__)
61
62			# Types
63			RedisCtx = NewType('RedisCtx', redis.Redis)
64
65
66			class OpenvasDB:
67			"""Class to connect to redis, to perform queries, and to move
68			from a KB to another."""
69
70			_db_address = None
71
72			@classmethod
73			def get_database_address(cls) -> Optional[str]:
74			if not cls._db_address:
75			settings = Openvas.get_settings()
76
77			cls._db_address = settings.get('db_address')
78
79			return cls._db_address
80
81			@classmethod
82			def create_context(
83			cls, dbnum: Optional[int] = 0, encoding: Optional[str] = 'latin-1'
84			) -> RedisCtx:
85			"""Connect to redis to the given database or to the default db 0 .
86
87			Arguments:
88			dbnum: The db number to connect to.
89			encoding: The encoding to be used to read and write.
90
91			Return a new redis context on success.
92			"""
93			tries = 5
94			while tries:
95			try:
96			ctx = redis.Redis(
97			unix_socket_path=cls.get_database_address(),
98			db=dbnum,
99			socket_timeout=SOCKET_TIMEOUT,
100			encoding=encoding,
101			decode_responses=True,
102			)
103			ctx.keys("test")
104			except (redis.exceptions.ConnectionError, FileNotFoundError) as err:
105			logger.debug(
106			'Redis connection lost: %s. Trying again in 5 seconds.', err
107			)
108			tries = tries - 1
109			time.sleep(5)
110			continue
111			break
112
113			if not tries:
114			logger.error('Redis Error: Not possible to connect to the kb.')
115			sys.exit(1)
116
117			return ctx
			0 ignored issues – show introduced 2021-03-23 10:44 UTC by Report Bug Copy Issue Report The variable `ctx` does not seem to be defined in case the `while` loop on line `94` is not entered. Are you sure this can never be the case? Loading history...
118
119			@classmethod
120			def find_database_by_pattern(
121			cls, pattern: str, max_database_index: int
122			) -> Tuple[Optional[RedisCtx], Optional[int]]:
123			"""Search a pattern inside all kbs up to max_database_index.
124
125			Returns the redis context for the db and its index as a tuple or
126			None, None if the db with the pattern couldn't be found.
127			"""
128			for i in range(0, max_database_index):
129			ctx = cls.create_context(i)
130			if ctx.keys(pattern):
131			return (ctx, i)
132
133			return (None, None)
134
135			@staticmethod
136			def select_database(ctx: RedisCtx, kbindex: str):
137			"""Use an existent redis connection and select a redis kb.
138
139			Arguments:
140			ctx: Redis context to use.
141			kbindex: The new kb to select
142			"""
143			if not ctx:
144			raise RequiredArgument('select_database', 'ctx')
145			if not kbindex:
146			raise RequiredArgument('select_database', 'kbindex')
147
148			ctx.execute_command('SELECT ' + str(kbindex))
149
150			@staticmethod
151			def get_list_item(
152			ctx: RedisCtx,
153			name: str,
154			start: Optional[int] = LIST_FIRST_POS,
155			end: Optional[int] = LIST_LAST_POS,
156			) -> Optional[list]:
157			"""Returns the specified elements from `start` to `end` of the
158			list stored as `name`.
159
160			Arguments:
161			ctx: Redis context to use.
162			name: key name of a list.
163			start: first range element to get.
164			end: last range element to get.
165
166			Return List specified elements in the key.
167			"""
168			if not ctx:
169			raise RequiredArgument('get_list_item', 'ctx')
170			if not name:
171			raise RequiredArgument('get_list_item', 'name')
172
173			return ctx.lrange(name, start, end)
174
175			@staticmethod
176			def get_last_list_item(ctx: RedisCtx, name: str) -> str:
177			if not ctx:
178			raise RequiredArgument('get_last_list_item', 'ctx')
179			if not name:
180			raise RequiredArgument('get_last_list_item', 'name')
181
182			return ctx.rpop(name)
183
184			@staticmethod
185			def pop_list_items(ctx: RedisCtx, name: str) -> List[str]:
186			if not ctx:
187			raise RequiredArgument('pop_list_items', 'ctx')
188			if not name:
189			raise RequiredArgument('pop_list_items', 'name')
190
191			pipe = ctx.pipeline()
192			pipe.lrange(name, LIST_FIRST_POS, LIST_LAST_POS)
193			pipe.delete(name)
194			results, redis_return_code = pipe.execute()
195
196			# The results are left-pushed. To preserver the order
197			# the result list must be reversed.
198			if redis_return_code:
199			results.reverse()
200			else:
201			results = []
202
203			return results
204
205			@staticmethod
206			def get_key_count(ctx: RedisCtx, pattern: Optional[str] = None) -> int:
207			"""Get the number of keys matching with the pattern.
208
209			Arguments:
210			ctx: Redis context to use.
211			pattern: pattern used as filter.
212			"""
213			if not pattern:
214			pattern = "*"
215
216			if not ctx:
217			raise RequiredArgument('get_key_count', 'ctx')
218
219			return len(ctx.keys(pattern))
220
221			@staticmethod
222			def remove_list_item(ctx: RedisCtx, key: str, value: str):
223			"""Remove item from the key list.
224
225			Arguments:
226			ctx: Redis context to use.
227			key: key name of a list.
228			value: Value to be removed from the key.
229			"""
230			if not ctx:
231			raise RequiredArgument('remove_list_item ', 'ctx')
232			if not key:
233			raise RequiredArgument('remove_list_item', 'key')
234			if not value:
235			raise RequiredArgument('remove_list_item ', 'value')
236
237			ctx.lrem(key, count=LIST_ALL, value=value)
238
239			@staticmethod
240			def get_single_item(
241			ctx: RedisCtx,
242			name: str,
243			index: Optional[int] = LIST_FIRST_POS,
244			) -> Optional[str]:
245			"""Get a single KB element.
246
247			Arguments:
248			ctx: Redis context to use.
249			name: key name of a list.
250			index: index of the element to be return.
251			Defaults to the first element in the list.
252
253			Return the first element of the list or None if the name couldn't be
254			found.
255			"""
256			if not ctx:
257			raise RequiredArgument('get_single_item', 'ctx')
258			if not name:
259			raise RequiredArgument('get_single_item', 'name')
260
261			return ctx.lindex(name, index)
262
263			@staticmethod
264			def add_single_list(ctx: RedisCtx, name: str, values: Iterable):
265			"""Add a single KB element with one or more values.
266			The values can be repeated. If the key already exists will
267			be removed an completely replaced.
268
269			Arguments:
270			ctx: Redis context to use.
271			name: key name of a list.
272			value: Elements to add to the key.
273			"""
274			if not ctx:
275			raise RequiredArgument('add_single_list', 'ctx')
276			if not name:
277			raise RequiredArgument('add_single_list', 'name')
278			if not values:
279			raise RequiredArgument('add_single_list', 'value')
280
281			pipe = ctx.pipeline()
282			pipe.delete(name)
283			pipe.rpush(name, *values)
284			pipe.execute()
285
286			@staticmethod
287			def add_single_item(ctx: RedisCtx, name: str, values: Iterable):
288			"""Add a single KB element with one or more values. Don't add
289			duplicated values during this operation, but if the the same
290			values already exists under the key, this will not be overwritten.
291
292			Arguments:
293			ctx: Redis context to use.
294			name: key name of a list.
295			value: Elements to add to the key.
296			"""
297			if not ctx:
298			raise RequiredArgument('add_single_item', 'ctx')
299			if not name:
300			raise RequiredArgument('add_single_item', 'name')
301			if not values:
302			raise RequiredArgument('add_single_item', 'value')
303
304			ctx.rpush(name, *set(values))
305
306			@staticmethod
307			def set_single_item(ctx: RedisCtx, name: str, value: Iterable):
308			"""Set (replace) a single KB element. If the same key exists
309			in the kb, it is completed removed. Values added are unique.
310
311			Arguments:
312			ctx: Redis context to use.
313			name: key name of a list.
314			value: New elements to add to the key.
315			"""
316			if not ctx:
317			raise RequiredArgument('set_single_item', 'ctx')
318			if not name:
319			raise RequiredArgument('set_single_item', 'name')
320			if not value:
321			raise RequiredArgument('set_single_item', 'value')
322
323			pipe = ctx.pipeline()
324			pipe.delete(name)
325			pipe.rpush(name, *set(value))
326			pipe.execute()
327
328			@staticmethod
329			def get_pattern(ctx: RedisCtx, pattern: str) -> List:
330			"""Get all items stored under a given pattern.
331
332			Arguments:
333			ctx: Redis context to use.
334			pattern: key pattern to match.
335
336			Return a list with the elements under the matched key.
337			"""
338			if not ctx:
339			raise RequiredArgument('get_pattern', 'ctx')
340			if not pattern:
341			raise RequiredArgument('get_pattern', 'pattern')
342
343			items = ctx.keys(pattern)
344
345			elem_list = []
346			for item in items:
347			elem_list.append(
348			[
349			item,
350			ctx.lrange(item, start=LIST_FIRST_POS, end=LIST_LAST_POS),
351			]
352			)
353			return elem_list
354
355			@classmethod
356			def get_keys_by_pattern(cls, ctx: RedisCtx, pattern: str) -> List[str]:
357			"""Get all items with index 'index', stored under
358			a given pattern.
359
360			Arguments:
361			ctx: Redis context to use.
362			pattern: key pattern to match.
363
364			Return a sorted list with the elements under the matched key
365			"""
366			if not ctx:
367			raise RequiredArgument('get_elem_pattern_by_index', 'ctx')
368			if not pattern:
369			raise RequiredArgument('get_elem_pattern_by_index', 'pattern')
370
371			return sorted(ctx.keys(pattern))
372
373			@classmethod
374			def get_filenames_and_oids(
375			cls,
376			ctx: RedisCtx,
377			) -> Iterable[Tuple[str, str]]:
378			"""Get all items with index 'index', stored under
379			a given pattern.
380
381			Arguments:
382			ctx: Redis context to use.
383
384			Return an iterable where each single tuple contains the filename
385			as first element and the oid as the second one.
386			"""
387			if not ctx:
388			raise RequiredArgument('get_filenames_and_oids', 'ctx')
389
390			items = cls.get_keys_by_pattern(ctx, 'nvt:*')
391
392			return ((ctx.lindex(item, 0), item[4:]) for item in items)
			0 ignored issues – show Comprehensibility Best Practice introduced 2020-03-26 16:28 UTC by Report Bug Copy Issue Report The variable `item` does not seem to be defined. Loading history...
393
394
395			class BaseDB:
396			def __init__(self, kbindex: int, ctx: Optional[RedisCtx] = None):
397			if ctx is None:
398			self.ctx = OpenvasDB.create_context(kbindex)
399			else:
400			self.ctx = ctx
401
402			self.index = kbindex
403
404			def flush(self):
405			"""Flush the database"""
406			self.ctx.flushdb()
407
408
409			class BaseKbDB(BaseDB):
410			def _add_single_item(
411			self, name: str, values: Iterable, utf8_enc: Optional[bool] = False
412			):
413			"""Changing the encoding format of an existing redis context
414			is not possible. Therefore a new temporary redis context is
415			created to store key-values encoded with utf-8."""
416			if utf8_enc:
417			ctx = OpenvasDB.create_context(self.index, encoding='utf-8')
418			OpenvasDB.add_single_item(ctx, name, values)
419			else:
420			OpenvasDB.add_single_item(self.ctx, name, values)
421
422			def _set_single_item(self, name: str, value: Iterable):
423			"""Set (replace) a single KB element.
424
425			Arguments:
426			name: key name of a list.
427			value: New elements to add to the key.
428			"""
429			OpenvasDB.set_single_item(self.ctx, name, value)
430
431			def _get_single_item(self, name: str) -> Optional[str]:
432			"""Get a single KB element.
433
434			Arguments:
435			name: key name of a list.
436			"""
437			return OpenvasDB.get_single_item(self.ctx, name)
438
439			def _get_list_item(
440			self,
441			name: str,
442			) -> Optional[List]:
443			"""Returns the specified elements from `start` to `end` of the
444			list stored as `name`.
445
446			Arguments:
447			name: key name of a list.
448
449			Return List specified elements in the key.
450			"""
451			return OpenvasDB.get_list_item(self.ctx, name)
452
453			def _pop_list_items(self, name: str) -> List:
454			return OpenvasDB.pop_list_items(self.ctx, name)
455
456			def _remove_list_item(self, key: str, value: str):
457			"""Remove item from the key list.
458
459			Arguments:
460			key: key name of a list.
461			value: Value to be removed from the key.
462			"""
463			OpenvasDB.remove_list_item(self.ctx, key, value)
464
465			def get_result(self) -> Optional[str]:
466			"""Get and remove the oldest result from the list.
467
468			Return the oldest scan results
469			"""
470			return self._pop_list_items("internal/results")
471
472			def get_status(self, openvas_scan_id: str) -> Optional[str]:
473			"""Return the status of the host scan"""
474			return self._get_single_item('internal/{}'.format(openvas_scan_id))
475
476			def __repr__(self):
477			return '<{} index={}>'.format(self.__class__.__name__, self.index)
478
479
480			class ScanDB(BaseKbDB):
481			"""Database for a scanning a single host"""
482
483			def select(self, kbindex: int) -> "ScanDB":
484			"""Select a redis kb.
485
486			Arguments:
487			kbindex: The new kb to select
488			"""
489			OpenvasDB.select_database(self.ctx, kbindex)
490			self.index = kbindex
491			return self
492
493
494			class KbDB(BaseKbDB):
495			def get_scan_databases(self) -> Iterator[ScanDB]:
496			"""Returns an iterator yielding corresponding ScanDBs
497
498			The returned Iterator can't be converted to an Iterable like a List.
499			Each yielded ScanDB must be used independently in a for loop. If the
500			Iterator gets converted into an Iterable all returned ScanDBs will use
501			the same redis context pointing to the same redis database.
502			"""
503			dbs = self._get_list_item('internal/dbindex')
504			scan_db = ScanDB(self.index)
505			for kbindex in dbs:
506			if kbindex == self.index:
507			continue
508
509			yield scan_db.select(kbindex)
510
511			def add_scan_id(self, scan_id: str):
512			self._add_single_item('internal/{}'.format(scan_id), ['new'])
513			self._add_single_item('internal/scanid', [scan_id])
514
515			def add_scan_preferences(self, openvas_scan_id: str, preferences: Iterable):
516			self._add_single_item(
517			'internal/{}/scanprefs'.format(openvas_scan_id), preferences
518			)
519
520			def add_credentials_to_scan_preferences(
521			self, openvas_scan_id: str, preferences: Iterable
522			):
523			"""Force the usage of the utf-8 encoding, since some credentials
524			contain special chars not supported by latin-1 encoding."""
525			self._add_single_item(
526			'internal/{}/scanprefs'.format(openvas_scan_id),
527			preferences,
528			utf8_enc=True,
529			)
530
531			def add_scan_process_id(self, pid: int):
532			self._add_single_item('internal/ovas_pid', [pid])
533
534			def get_scan_process_id(self) -> Optional[str]:
535			return self._get_single_item('internal/ovas_pid')
536
537			def remove_scan_database(self, scan_db: ScanDB):
538			self._remove_list_item('internal/dbindex', scan_db.index)
539
540			def target_is_finished(self, scan_id: str) -> bool:
541			"""Check if a target has finished."""
542
543			status = self._get_single_item('internal/{}'.format(scan_id))
544
545			if status is None:
546			logger.error(
547			"%s: Target set as finished because redis returned None as "
548			"scanner status.",
549			scan_id,
550			)
551
552			return status == 'finished' or status is None
553
554			def stop_scan(self, openvas_scan_id: str):
555			self._set_single_item(
556			'internal/{}'.format(openvas_scan_id), ['stop_all']
557			)
558
559			def scan_is_stopped(self, scan_id: str) -> bool:
560			"""Check if the scan should be stopped"""
561			status = self._get_single_item('internal/%s' % scan_id)
562			return status == 'stop_all'
563
564			def get_scan_status(self) -> List:
565			"""Get and remove the oldest host scan status from the list.
566
567			Return a string which represents the host scan status.
568			"""
569			return self._pop_list_items("internal/status")
570
571
572			class MainDB(BaseDB):
573			"""Main Database"""
574
575			DEFAULT_INDEX = 0
576
577			def __init__(self, ctx=None):
578			super().__init__(self.DEFAULT_INDEX, ctx)
579
580			self._max_dbindex = None
581
582			@property
583			def max_database_index(self):
584			"""Set the number of databases have been configured into kbr struct."""
585			if self._max_dbindex is None:
586			resp = self.ctx.config_get('databases')
587
588			if len(resp) == 1:
589			self._max_dbindex = int(resp.get('databases'))
590			else:
591			raise OspdOpenvasError(
592			'Redis Error: Not possible to get max_dbindex.'
593			) from None
594
595			return self._max_dbindex
596
597			def try_database(self, index: int) -> bool:
598			"""Check if a redis db is already in use. If not, set it
599			as in use and return.
600
601			Arguments:
602			ctx: Redis object connected to the kb with the
603			DBINDEX_NAME key.
604			index: Number intended to be used.
605
606			Return True if it is possible to use the db. False if the given db
607			number is already in use.
608			"""
609			_in_use = 1
610			try:
611			resp = self.ctx.hsetnx(DBINDEX_NAME, index, _in_use)
612			except:
613			raise OspdOpenvasError(
614			'Redis Error: Not possible to set %s.' % DBINDEX_NAME
615			) from None
616
617			return resp == 1
618
619			def get_new_kb_database(self) -> Optional[KbDB]:
620			"""Return a new kb db to an empty kb."""
621			for index in range(1, self.max_database_index):
622			if self.try_database(index):
623			kbdb = KbDB(index)
624			kbdb.flush()
625			return kbdb
626
627			return None
628
629			def find_kb_database_by_scan_id(
630			self, scan_id: str
631			) -> Tuple[Optional[str], Optional["KbDB"]]:
632			"""Find a kb db by via a scan id"""
633			for index in range(1, self.max_database_index):
634			ctx = OpenvasDB.create_context(index)
635			if OpenvasDB.get_key_count(ctx, 'internal/{}'.format(scan_id)):
636			return KbDB(index, ctx)
637
638			return None
639
640			def release_database(self, database: BaseDB):
641			self.release_database_by_index(database.index)
642			database.flush()
643
644			def release_database_by_index(self, index: int):
645			self.ctx.hdel(DBINDEX_NAME, index)
646
647			def release(self):
648			self.release_database(self)
649

greenbone / ospd-openvas

ospd_openvas.db.KbDB.add_scan_id() A last analyzed 2021-07-12 05:11 UTC

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like

ospd_openvas.db.KbDB.add_scan_id() A
last analyzed 2021-07-12 05:11 UTC