From 4a92b863f9eab18b9616f3596d6e34d489631892 Mon Sep 17 00:00:00 2001
From: luigi Giacobbe <lgiacobbe@gmail.com>
Date: Tue, 5 Apr 2022 00:35:54 +0200
Subject: [PATCH] generic device functions added

---
 pytos/securetrack/helpers.py | 84 ++++++++++++++++++++++++++++++++++++
 requirements.txt             |  2 +
 2 files changed, 86 insertions(+)

diff --git a/pytos/securetrack/helpers.py b/pytos/securetrack/helpers.py
index 5d66482..6394091 100644
--- a/pytos/securetrack/helpers.py
+++ b/pytos/securetrack/helpers.py
@@ -6,6 +6,8 @@
 import logging
 import base64
 import multiprocessing.pool
+import json
+import requests
 import xml.etree.ElementTree as ET
 
 from pytos.common.definitions.Url_Params_Builder import URLParamBuilderDict
@@ -299,6 +301,88 @@ def get_generic_device_by_name(self, generic_device_name, domain_id=DEFAULT_DOMA
             logger.error(msg)
             raise ValueError(msg)
 
+    def new_generic_device(self, name: str, config_file: str,  domain_id=1, update_topo=True):
+        """
+        It creates a new generic device on the specified domain.
+
+        :param name: the generic device name
+        :param config_file: the configuration file to upload to ST
+        :param domain_id: the domain id where the device needs to be created
+        :param update_topo: update the topology after the creation
+        """
+        gen_device = json.dumps({"generic_device": {"name": name, "customer_id": domain_id}})
+        multi_part_form_dict = {"device_data": ('data.txt', gen_device, 'application/json'),
+                                "configuration_file": (
+                                config_file, open(config_file, 'rb'), "application/octet-stream"),
+                                'update_topology': (None, update_topo, 'text/plain')
+                                }
+
+        try:
+            response = self.post_uri("/securetrack/api/generic_devices/",
+                                     multi_part_form_params=multi_part_form_dict, expected_status_codes=201)
+            logger.info("creation of generic device {} is successful.".format(name))
+        except RequestException:
+            message = "Failed to create the new device with name {}".format(name)
+            logger.critical(message)
+            raise IOError(message)
+        except REST_Bad_Request_Error as http_exception:
+            message = "Failed to create the new device with name {}, got error: '{}'.".format(name, http_exception)
+            logger.critical(message)
+            raise ValueError(message)
+
+    def update_generic_device(self, device_id: int, config_file: str, update_topology=False):
+        """
+        The update_generic_device function updates a generic device with the given configuration file.
+
+        :param self: Used to Reference the class instance.
+        :param device_id: Used to Pass in the id of the device that we want to update.
+        :param config_file: Used to Specify the path to the configuration file for the device.
+        :param update_topology: Used to Tell the update_generic_device function to not update the topology.
+        :return: :.
+        """
+
+        openBin = {'configuration_file': (config_file, open(config_file, 'rb'), 'application/octet-stream')}
+        try:
+            url = 'https://' + self._real_hostname + '/securetrack/api/generic_devices/{}'.format(device_id)
+
+            t = requests.put(url, headers={'Content-Type': 'multipart/form-data'},
+                             files=openBin, data={"update_topology": update_topology}, verify=False,
+                             auth=(self.login_data['username'], self.login_data['password']))
+            if t.status_code != 204:
+                message = "Failed to update the device with id {}.  status code is {}".format(device_id, t.status_code)
+                logger.critical(message)
+                raise ValueError(message)
+
+            logger.info("Upload successful.")
+
+        except RequestException:
+            message = "Failed to update the new device with id {}".format(device_id)
+            logger.critical(message)
+            raise IOError(message)
+        except REST_Bad_Request_Error as http_exception:
+            message = "Failed to create the new device with id {}, got error: '{}'.".format(device_id, http_exception)
+            logger.critical(message)
+            raise ValueError(message)
+
+    def delete_generic_device(self, device_id: int):
+        """
+        It deletes a generic device from SecureTrack map
+
+        :parameter device_id: Used to Pass in the id of the device that we want to delete
+        """
+        logger.info("Deleting generic device with ID {}.".format(device_id))
+        url = "/securetrack/api/generic_devices/{}?update_topology=false".format(device_id)
+        try:
+            self.delete_uri(url, expected_status_codes=[200, 204])
+        except REST_Not_Found_Error:
+            message = "Generic device with ID {} doesn't exist.".format(device_id)
+            logger.critical(message)
+            raise ValueError(message)
+        except (RequestException, REST_HTTP_Exception):
+            message = "Failed to delete generic device with ID {}.".format(device_id)
+            logger.critical(message)
+            raise IOError(message)
+
     def add_offline_device(self, name, vendor, model, domain_id=None, domain_name="Default", topology="true",
                            offline="true"):
         """Add an offline device to SecureTrack.
diff --git a/requirements.txt b/requirements.txt
index 180ab04..bc0d4e1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,3 +3,5 @@ wheel==0.24.0
 lxml>=4.6.2
 pyinotify==0.9.6
 netifaces==0.10.9
+
+setuptools~=60.2.0
\ No newline at end of file