ospd_openvas.db.KbDB.add_credentials_to_scan_preferences() - Code Metrics - Inspection of "MQTT Database for results, Option to enable MQTT" - greenbone/ospd-openvas - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Pull Request — master (#422)

unknown

created 2021-05-19 08:43 UTC

KbDB.add_credentials_to_scan_preferences() A

↳ Parent: ospd_openvas.db

Complexity

Conditions

Size

Total Lines	9
Code Lines	6

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
eloc	6
nop	3
dl	0
loc	9
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
from queue import Queue

import redis
import paho.mqtt.client as mqtt

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 MQTTDB:
    def __init__(self, host):
        self.results = Queue()
        self.client = mqtt.Client(
            client_id="result_handler",
            protocol=mqtt.MQTTv5,
            userdata=self.results,
        )
        self.client.connect(host)
        self.client.on_message = self.on_message
        self.client.subscribe("scanner/results")
        self.client.loop_start()

    @staticmethod
    def on_message(client, userdata, msg):
        if msg.topic == "scanner/results":
            userdata.put(str(msg.payload.decode("utf-8")), False)

    def get_result(self):
        """Get all result from the result Queue"""
        ret = []
        while not self.results.empty():
            ret.append(self.results.get(False))
        return ret


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))

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

greenbone / ospd-openvas

Pull Request — master (#422)

KbDB.add_credentials_to_scan_preferences() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like