Get memory usage

CHANGELOG.md

@@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Extend osp to accept target options. [#194](https://github.com/greenbone/ospd/pull/194)
 - Accept reverse_lookup_only and reverse_lookup_unify target's options. [#195](https://github.com/greenbone/ospd/pull/195)
 - Add 'total' and 'sent' attributes to <vts> element for <get_vts> cmd response. [#206](https://github.com/greenbone/ospd/pull/206)
+- Add new get_memory_usage command. [#207](https://github.com/greenbone/ospd/pull/207)
 
 ### Changes
 - Modify __init__() method and use new syntax for super(). [#186](https://github.com/greenbone/ospd/pull/186)

Pipfile

@@ -8,6 +8,7 @@ lxml = "*"
 paramiko = "*"
 defusedxml = "*"
 deprecated = "*"
+psutil = "*"
 
 [dev-packages]
 coverage = "*"

doc/OSP.xml

@@ -1326,6 +1326,88 @@ Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
       </response>
     </example>
   </command>
+
+  <command>
+    <name>get_memory_usage</name>
+    <summary>Return memory usage information of the osp daemon</summary>
+    <pattern>
+      <attrib>
+        <name>unit</name>
+        <summary>Size unit for the memory. b for bytes, kb for kilobytes and mb for megabytes.</summary>
+        <type>text</type>
+      </attrib>
+    </pattern>
+    <response>
+      <pattern>
+        <attrib>
+          <name>status</name>
+          <type>status</type>
+          <required>1</required>
+        </attrib>
+        <attrib>
+          <name>status_text</name>
+          <type>text</type>
+          <required>1</required>
+        </attrib>
+      </pattern>
+      <ele>
+        <name>processes</name>
+        <summary>List of running processes</summary>
+        <ele>
+          <name>process</name>
+          <summary>Single running processes</summary>
+          <pattern>
+            <attrib>
+              <name>name</name>
+              <summary>Name of the process</summary>
+              <type>string</type>
+            </attrib>
+            <attrib>
+              <name>pid</name>
+              <summary>Process ID</summary>
+              <type>int</type>
+            </attrib>
+          </pattern>
+          <ele>
+            <name>rss</name>
+            <description>Resident Set Size. Non-swapped physical memory of the process</description>
+            <type>int</type>
+          </ele>
+          <ele>
+            <name>vms</name>
+            <description>Virtual Memory Size. Total amount of virtual memory used by the process.</description>
+            <type>int</type>
+          </ele>
+          <ele>
+            <name>shared</name>
+            <description>Memory shared with other processes</description>
+            <type>int</type>
+          </ele>
+        </ele>
+      </ele>
+      </pattern>
+    </response>
+    <example>
+      <request>
+        <get_memory_usage unit="kb"/>
+      </request>
+      <response>
+        <get_memory_usage status="200" status_text="OK">
+          <processes>
+            <process name="MainProcess" pid="12345">
+              <rss>127182</rss>
+              <vss>239616</vss>
+              <shared>135168</shared>
+            </process>
+            <process name="Process-1" pid="23456">
+            ...
+            </process>
+          </processes>
+        </get_memory_usage>
+      </response>
+    </example>
+  </command>
+
   <parameter_type>
     <name>integer</name>
     <summary>An integer value</summary>

ospd/command/command.py

@@ -16,13 +16,17 @@
 # along with this program; if not, write to the Free Software
 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 
+import multiprocessing
 import re
 import subprocess
 
+from decimal import Decimal
 from typing import Optional, Dict, Any, Union, Iterator
 
 from xml.etree.ElementTree import Element, SubElement
 
+import psutil
+
 from ospd.errors import OspdCommandError
 from ospd.misc import valid_uuid, create_process
 from ospd.network import target_str_to_list
@@ -354,7 +358,9 @@ def handle_xml(self, xml: Element) -> bytes:
         # Don't send response until the scan is stopped.
         try:
             self._daemon.scan_processes[scan_id].join()
-            exitcode = self._daemon.scan_processes[scan_id].exitcode
+            exitcode = self._daemon.scan_processes[  # pylint: disable=unused-variable
+                scan_id
+            ].exitcode
         except KeyError:
             pass
 
@@ -522,3 +528,84 @@ def handle_xml(self, xml: Element) -> bytes:
         id_.text = scan_id
 
         return simple_response_str('start_scan', 200, 'OK', id_)
+
+
+class GetMemoryUsage(BaseCommand):
+
+    name = "get_memory_usage"
+    description = "print the memory consumption of all processes"
+    attributes = {
+        'unit': 'Unit for displaying memory consumption (b = bytes, '
+        'kb = kilobytes, mb = megabytes). Defaults to b.'
+    }
+
+    @staticmethod
+    def _get_memory(value: int, unit: str = None) -> str:
+        if not unit:
+            return str(value)
+
+        unit = unit.lower()
+
+        if unit == 'kb':
+            return str(Decimal(value) / 1024)
+
+        if unit == 'mb':
+            return str(Decimal(value) / (1024 * 1024))
+
+        return str(value)
+
+    @staticmethod
+    def _create_process_element(name: str, pid: int):
+        process_element = Element('process')
+        process_element.set('name', name)
+        process_element.set('pid', str(pid))
+
+        return process_element
+
+    @classmethod
+    def _add_memory_info(
+        cls, process_element: Element, pid: int, unit: str = None
+    ):
+        try:
+            ps_process = psutil.Process(pid)
+        except psutil.NoSuchProcess:
+            return
+
+        memory = ps_process.memory_info()
+
+        rss_element = Element('rss')
+        rss_element.text = cls._get_memory(memory.rss, unit)
+
+        process_element.append(rss_element)
+
+        vms_element = Element('vms')
+        vms_element.text = cls._get_memory(memory.vms, unit)
+
+        process_element.append(vms_element)
+
+        shared_element = Element('shared')
+        shared_element.text = cls._get_memory(memory.shared, unit)
+
+        process_element.append(shared_element)
+
+    def handle_xml(self, xml: Element) -> bytes:
+        processes_element = Element('processes')
+        unit = xml.get('unit')
+
+        current_process = multiprocessing.current_process()
+        process_element = self._create_process_element(
+            current_process.name, current_process.pid
+        )
+
+        self._add_memory_info(process_element, current_process.pid, unit)
+
+        processes_element.append(process_element)
+
+        for proc in multiprocessing.active_children():
+            process_element = self._create_process_element(proc.name, proc.pid)
+
+            self._add_memory_info(process_element, proc.pid, unit)
+
+            processes_element.append(process_element)
+
+        return simple_response_str('get_memory', 200, 'OK', processes_element)

setup.py

@@ -81,7 +81,7 @@
     # project is installed. For an analysis of "install_requires" vs pip's
     # requirements files see:
     # https://packaging.python.org/en/latest/technical.html#install-requires-vs-requirements-files
-    install_requires=['paramiko', 'defusedxml', 'lxml', 'deprecated'],
+    install_requires=['paramiko', 'defusedxml', 'lxml', 'deprecated', 'psutil'],
     # You can just specify the packages manually here if your project is
     # simple. Or you can use find_packages().
     packages=find_packages(exclude=['tests']),