Merge pull request #44 from bjoernricks/connection-updates

Connection updates

Author: jjnicola

CHANGELOG.md

@@ -10,7 +10,7 @@ Some notable changes are:
   having to select the connection when creating the gmp object.
 * Support for different protocols and versions has been added. Currently
   supported protocols are OSP v1 and GMP v7.
-* Full API documentation is available.
+* Full API documentation is available at https://python-gvm.readthedocs.io/en/latest/.
 * Possible arguments to protocol methods are documented.
 * Arguments should be passed as keywords
 

gvm/connections.py

@@ -66,39 +66,75 @@ def _feed_xml(self, data):
                            "read {0}".format(data), e)
 
 
-class GvmConnection:
+class GvmConnection(XmlReader):
     """
     Base class for establishing a connection to a remote server daemon.
+
+    Arguments:
+        timeout (int, optional): Timeout in seconds for the connection.
     """
 
     def __init__(self, timeout=DEFAULT_TIMEOUT):
-        """
-          Arguments:
-            socket -- A socket
-        """
         self._socket = None
         self._timeout = timeout
 
+    def _read(self):
+        return self._socket.recv(BUF_SIZE)
+
     def connect(self):
-        """Establish a connection to gvmd
+        """Establish a connection to a remote server
         """
         raise NotImplementedError
 
     def send(self, data):
-        """Send data to gvmd
+        """Send data to the connected remote server
+
+        Arguments:
+            data (str or bytes): Data to be send to the server. Either utf-8
+                encoded string or bytes.
         """
         if isinstance(data, str):
-            self._socket.send(data.encode())
+            self._socket.sendall(data.encode())
         else:
-            self._socket.send(data)
+            self._socket.sendall(data)
 
     def read(self):
-        """Read data from gvmd
+        """Read data from the remote server
+
+        Returns:
+            str: data as utf-8 encoded string
         """
-        raise NotImplementedError
+        response = ''
+
+        self._start_xml()
+
+        now = time.time()
+
+        break_timeout = now + self._timeout
+
+        while True:
+            data = self._read()
+
+            if not data:
+                # Connection was closed by server
+                raise GvmError('Remote closed the connection')
+
+            self._feed_xml(data)
+
+            response += data.decode('utf-8', errors='ignore')
+
+            if self._is_end_xml():
+                break
+
+            now = time.time()
+
+            if now > break_timeout:
+                raise GvmError('Timeout while reading the response')
+
+        return response
 
     def disconnect(self):
-        """Close the connection to gvmd
+        """Disconnect and close the connection to the remote server
         """
         try:
             if self._socket is not None:
@@ -107,9 +143,17 @@ def disconnect(self):
             logger.debug('Connection closing error: %s', e)
 
 
-class SSHConnection(GvmConnection, XmlReader):
+class SSHConnection(GvmConnection):
     """
     SSH Class to connect, read and write from GVM via SSH
+
+    Arguments:
+        timeout (int, optional): Timeout in seconds for the connection.
+        hostname (str, optional): DNS name or IP address of the remote server.
+            Default is 127.0.0.1.
+        port (int, optional): Port of the remote SSH server.
+        username (str, optional): Username to use for SSH login.
+        password (str, optional): Passwort to use for SSH login.
     """
 
     def __init__(self, timeout=DEFAULT_TIMEOUT, hostname='127.0.0.1', port=22,
@@ -143,6 +187,9 @@ def _send_in_chunks(self, data, chunk_size):
         return sent_bytes
 
     def connect(self):
+        """
+        Connect to the SSH server and authenticate to it
+        """
         self._socket = paramiko.SSHClient()
         self._socket.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 
@@ -160,35 +207,16 @@ def connect(self):
 
         except (paramiko.BadHostKeyException,
                 paramiko.AuthenticationException,
-                paramiko.SSHException, OSError) as e:
-            logger.debug('SSH Connection failed: %s', e)
-            raise
-
-    def read(self):
-        response = ''
-
-        self._start_xml()
-
-        while True:
-            data = self._stdout.channel.recv(BUF_SIZE)
-            # Connection was closed by server
-            if not data:
-                break
+                paramiko.SSHException,
+                ) as e:
+            raise GvmError('SSH Connection failed', e)
 
-            self._feed_xml(data)
-
-            response += data.decode('utf-8', errors='ignore')
-
-            if self._is_end_xml():
-                break
-
-        return response
+    def _read(self):
+        return self._stdout.channel.recv(BUF_SIZE)
 
     def send(self, data):
-        logger.debug('SSH:send(): %s', data)
         if len(data) > MAX_SSH_DATA_LENGTH:
-            sent_bytes = self._send_in_chunks(data, MAX_SSH_DATA_LENGTH)
-            logger.debug("SSH: %s bytes sent.", sent_bytes)
+            self._send_in_chunks(data, MAX_SSH_DATA_LENGTH)
         else:
             self._stdin.channel.send(data)
 
@@ -197,10 +225,28 @@ class TLSConnection(GvmConnection):
     """
     TLS class to connect, read and write from a remote GVM daemon via TLS
     secured socket.
+
+    Arguments:
+        timeout (int, optional): Timeout in seconds for the connection.
+        hostname (str, optional): DNS name or IP address of the remote TLS
+            server.
+        port (str, optional): Port for the TLS connection. Default is 9390.
+        certfile (str, optional): Path to PEM encoded certificate file. See
+            `python certificates`_ for details.
+        cafile (str, optional): Path to PEM encoded CA file. See
+            `python certificates`_ for details.
+        keyfile (str, optional): Path to PEM encoded private key. See
+            `python certificates`_ for details.
+        password (str, optional): Password for the private key. If the password
+            argument is not specified and a password is required it will be
+            interactively prompt the user for a password.
+
+    .. _python certificates:
+        https://docs.python.org/3.5/library/ssl.html#certificates
     """
 
     def __init__(self, certfile=None, cafile=None, keyfile=None,
-                 hostname='127.0.0.1', port=DEFAULT_GVM_PORT,
+                 hostname='127.0.0.1', port=DEFAULT_GVM_PORT, password=None,
                  timeout=DEFAULT_TIMEOUT):
         super().__init__(timeout=timeout)
 
@@ -209,45 +255,41 @@ def __init__(self, certfile=None, cafile=None, keyfile=None,
         self.certfile = certfile
         self.cafile = cafile
         self.keyfile = keyfile
+        self.password = password
 
     def _new_socket(self):
+        transport_socket = socketlib.socket(socketlib.AF_INET,
+                                            socketlib.SOCK_STREAM)
+
         if self.certfile and self.cafile and self.keyfile:
             context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH,
                                                  cafile=self.cafile)
             context.check_hostname = False
             context.load_cert_chain(
-                certfile=self.certfile, keyfile=self.keyfile)
-            new_socket = socketlib.socket(socketlib.AF_INET,
-                                          socketlib.SOCK_STREAM)
-            sock = context.wrap_socket(new_socket, server_side=False)
+                certfile=self.certfile, keyfile=self.keyfile,
+                password=self.password)
+            sock = context.wrap_socket(transport_socket, server_side=False)
         else:
             context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
-            sock = context.wrap_socket(socketlib.socket(socketlib.AF_INET))
-        return sock
+            sock = context.wrap_socket(transport_socket)
 
+        sock.settimeout(self._timeout)
+
+        return sock
 
     def connect(self):
         self._socket = self._new_socket()
-        self._socket.settimeout(self._timeout)
         self._socket.connect((self.hostname, int(self.port)))
 
-    def read(self):
-        response = ''
 
-        while True:
-            data = self._socket.read(BUF_SIZE)
-
-            response += data.decode('utf-8', errors='ignore')
-            if len(data) < BUF_SIZE:
-                break
-
-        return response
-
-
-class UnixSocketConnection(GvmConnection, XmlReader):
+class UnixSocketConnection(GvmConnection):
     """
     UNIX-Socket class to connect, read, write from a GVM server daemon via
     direct communicating UNIX-Socket
+
+    Arguments:
+        path (str, optional): Path to the socket.
+        timeout (int, optional): Timeout in seconds for the connection.
     """
 
     def __init__(self, path=DEFAULT_UNIX_SOCKET_PATH, timeout=DEFAULT_TIMEOUT,
@@ -265,34 +307,33 @@ def connect(self):
         self._socket.settimeout(self._timeout)
         self._socket.connect(self.path)
 
+
+class DebugConnection:
+
+    def __init__(self, connection):
+        self._connection = connection
+
     def read(self):
-        """Read from the UNIX socket
-        """
-        response = ''
+        data = self._connection.read()
 
-        break_timeout = time.time() + self.read_timeout
-        old_timeout = self._socket.gettimeout()
-        self._socket.settimeout(5)  # in seconds
+        logger.debug('Read %s characters. Data %s', len(data), data)
 
-        self._start_xml()
+        self.last_read_data = data
+        return data
 
-        while time.time() < break_timeout:
-            data = b''
+    def send(self, data):
+        self.last_send_data = data
 
-            try:
-                data = self._socket.recv(BUF_SIZE)
-            except (socketlib.timeout) as exception:
-                logger.debug('Warning: No data received '
-                             'from server: %s', exception)
-                continue
+        logger.debug('Sending %s characters. Data %s', len(data), data)
 
-            self._feed_xml(data)
+        return self._connection.send(data)
 
-            response += data.decode('utf-8', errors='ignore')
+    def connect(self):
+        logger.debug('Connecting')
 
-            if len(data) < BUF_SIZE:
-                if self._is_end_xml():
-                    break
+        return self._connection.connect()
 
-        self._socket.settimeout(old_timeout)
-        return response
+    def disconnect(self):
+        logger.debug('Disconnecting')
+
+        return self._connection.disconnect()