-
Notifications
You must be signed in to change notification settings - Fork 62
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Change: Rework and improve GvmConnection classes
The GvmConnection send method only accepts bytes now as only bytes should be send over a socket. Converting from and to string should be handled in the upper layers. SSHConnection got improved for better error handling. It leaked socket connections when errors occurred during receiving remote host information. The constructor got extended to be more flexible about print output, getting input and exiting. The new arguments allow also for much easier testing without mocking.
- Loading branch information
1 parent
c0e7d88
commit 94e3998
Showing
12 changed files
with
822 additions
and
757 deletions.
There are no files selected for viewing
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,32 @@ | ||
# SPDX-FileCopyrightText: 2024 Greenbone AG | ||
# | ||
# SPDX-License-Identifier: GPL-3.0-or-later | ||
|
||
from ._connection import DEFAULT_TIMEOUT, GvmConnection | ||
from ._debug import DebugConnection | ||
from ._ssh import ( | ||
DEFAULT_HOSTNAME, | ||
DEFAULT_KNOWN_HOSTS_FILE, | ||
DEFAULT_SSH_PASSWORD, | ||
DEFAULT_SSH_PORT, | ||
DEFAULT_SSH_USERNAME, | ||
SSHConnection, | ||
) | ||
from ._tls import DEFAULT_GVM_PORT, TLSConnection | ||
from ._unix import DEFAULT_UNIX_SOCKET_PATH, UnixSocketConnection | ||
|
||
__all__ = ( | ||
"DEFAULT_TIMEOUT", | ||
"DEFAULT_UNIX_SOCKET_PATH", | ||
"DEFAULT_GVM_PORT", | ||
"DEFAULT_HOSTNAME", | ||
"DEFAULT_KNOWN_HOSTS_FILE", | ||
"DEFAULT_SSH_PASSWORD", | ||
"DEFAULT_SSH_USERNAME", | ||
"DEFAULT_SSH_PORT", | ||
"DebugConnection", | ||
"GvmConnection", | ||
"SSHConnection", | ||
"TLSConnection", | ||
"UnixSocketConnection", | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,102 @@ | ||
# SPDX-FileCopyrightText: 2024 Greenbone AG | ||
# | ||
# SPDX-License-Identifier: GPL-3.0-or-later | ||
|
||
import logging | ||
import socket as socketlib | ||
from abc import ABC, abstractmethod | ||
from time import time | ||
from typing import Optional, Protocol, Union, runtime_checkable | ||
|
||
from gvm.errors import GvmError | ||
|
||
BUF_SIZE = 16 * 1024 | ||
|
||
DEFAULT_TIMEOUT = 60 # in seconds | ||
|
||
logger = logging.getLogger(__name__) | ||
|
||
|
||
@runtime_checkable | ||
class GvmConnection(Protocol): | ||
def connect(self) -> None: ... | ||
|
||
def disconnect(self) -> None: ... | ||
|
||
def send(self, data: bytes) -> None: ... | ||
|
||
def read(self) -> bytes: ... | ||
|
||
def finish_send(self): ... | ||
|
||
|
||
class AbstractGvmConnection(ABC): | ||
""" | ||
Base class for establishing a connection to a remote server daemon. | ||
Arguments: | ||
timeout: Timeout in seconds for the connection. None to | ||
wait indefinitely | ||
""" | ||
|
||
def __init__(self, timeout: Optional[Union[int, float]] = DEFAULT_TIMEOUT): | ||
self._socket: Optional[socketlib.SocketType] = None | ||
self._timeout = timeout if timeout is not None else DEFAULT_TIMEOUT | ||
|
||
def _read(self) -> bytes: | ||
if self._socket is None: | ||
raise GvmError("Socket is not connected") | ||
|
||
return self._socket.recv(BUF_SIZE) | ||
|
||
@abstractmethod | ||
def connect(self) -> None: | ||
"""Establish a connection to a remote server""" | ||
raise NotImplementedError | ||
|
||
def send(self, data: bytes) -> None: | ||
"""Send data to the connected remote server | ||
Arguments: | ||
data: Data to be send to the server. Either utf-8 encoded string or | ||
bytes. | ||
""" | ||
if self._socket is None: | ||
raise GvmError("Socket is not connected") | ||
|
||
self._socket.sendall(data) | ||
|
||
def read(self) -> bytes: | ||
"""Read data from the remote server | ||
Returns: | ||
str: data as utf-8 encoded string | ||
""" | ||
break_timeout = ( | ||
time() + self._timeout if self._timeout is not None else None | ||
) | ||
|
||
data = self._read() | ||
|
||
if not data: | ||
# Connection was closed by server | ||
raise GvmError("Remote closed the connection") | ||
|
||
if break_timeout and time() > break_timeout: | ||
raise GvmError("Timeout while reading the response") | ||
|
||
return data | ||
|
||
def disconnect(self) -> None: | ||
"""Disconnect and close the connection to the remote server""" | ||
try: | ||
if self._socket is not None: | ||
self._socket.close() | ||
except OSError as e: | ||
logger.debug("Connection closing error: %s", e) | ||
|
||
def finish_send(self): | ||
"""Indicate to the remote server you are done with sending data""" | ||
if self._socket is not None: | ||
# shutdown socket for sending. only allow reading data afterwards | ||
self._socket.shutdown(socketlib.SHUT_WR) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,71 @@ | ||
# SPDX-FileCopyrightText: 2024 Greenbone AG | ||
# | ||
# SPDX-License-Identifier: GPL-3.0-or-later | ||
|
||
import logging | ||
|
||
from ._connection import GvmConnection | ||
|
||
logger = logging.getLogger("gvm.connections.debug") | ||
|
||
|
||
class DebugConnection: | ||
"""Wrapper around a connection for debugging purposes | ||
Allows to debug the connection flow including send and read data. Internally | ||
it uses the python `logging`_ framework to create debug messages. Please | ||
take a look at `the logging tutorial | ||
<https://docs.python.org/3/howto/logging.html#logging-basic-tutorial>`_ | ||
for further details. | ||
Example: | ||
.. code-block:: python | ||
import logging | ||
logging.basicConfig(level=logging.DEBUG) | ||
socket_connection = UnixSocketConnection(path='/var/run/gvm.sock') | ||
connection = DebugConnection(socket_connection) | ||
gmp = Gmp(connection=connection) | ||
Arg: | ||
connection: GvmConnection to observe | ||
.. _logging: | ||
https://docs.python.org/3/library/logging.html | ||
""" | ||
|
||
def __init__(self, connection: GvmConnection): | ||
self._connection = connection | ||
|
||
def read(self) -> bytes: | ||
data = self._connection.read() | ||
|
||
logger.debug("Read %s characters. Data %r", len(data), data) | ||
|
||
self.last_read_data = data | ||
return data | ||
|
||
def send(self, data: bytes) -> None: | ||
self.last_send_data = data | ||
|
||
logger.debug("Sending %s characters. Data %r", len(data), data) | ||
|
||
return self._connection.send(data) | ||
|
||
def connect(self) -> None: | ||
logger.debug("Connecting") | ||
|
||
return self._connection.connect() | ||
|
||
def disconnect(self) -> None: | ||
logger.debug("Disconnecting") | ||
|
||
return self._connection.disconnect() | ||
|
||
def finish_send(self) -> None: | ||
logger.debug("Finish send") | ||
|
||
self._connection.finish_send() |
Oops, something went wrong.