Cleaned the code of the AttackMitigator component

LOG_LEVEL = logging.WARNING

# gRPC settings

GRPC_SERVICE_PORT = 10002  # TODO UPM FIXME

GRPC_SERVICE_PORT = 10002

GRPC_MAX_WORKERS = 10

GRPC_GRACE_PERIOD = 60

from __future__ import print_function

import logging

import grpc

import time

import json

from common.proto.l3_centralizedattackdetector_pb2 import Empty

from common.proto.l3_attackmitigator_pb2_grpc import L3AttackmitigatorServicer

from common.proto.context_pb2 import (

    Service,

    ServiceId,

    ServiceConfig,

    ServiceTypeEnum,

    ServiceStatusEnum,

    ServiceStatus,

    Context,

    ContextId,

    Uuid,

    Timestamp,

    ConfigRule,

    ConfigRule_Custom,

    ConfigActionEnum,

    Device,

    DeviceId,

    DeviceConfig,

    DeviceOperationalStatusEnum,

    DeviceDriverEnum,

    EndPoint,

    Link,

    LinkId,

    EndPoint,

    EndPointId,

    Topology,

    TopologyId,

)

from common.proto.context_pb2_grpc import ContextServiceStub

from common.proto.service_pb2_grpc import ServiceServiceStub

from datetime import datetime

import grpc

import time

import json

LOGGER = logging.getLogger(__name__)

CONTEXT_CHANNEL = "192.168.165.78:1010"

class l3_attackmitigatorServiceServicerImpl(L3AttackmitigatorServicer):

    def __init__(self):

        LOGGER.debug("Creating Servicer...")

        LOGGER.info("Creating Attack Mitigator Service")

        self.last_value = -1

        self.last_tag = 0

            "transport:destination-port": port_d,

            "forwarding-action": "DROP",

        }

        return value

    def GenerateContextId(self, context_id):

        uuid = Uuid()

        uuid.uuid = context_id

        context_id_obj.context_uuid.CopyFrom(uuid)

        return context_id_obj

    def GenerateServiceId(self, service_id):

        context_id.context_uuid.CopyFrom(uuid)

        service_id_obj.context_id.CopyFrom(context_id)

        service_id_obj.service_uuid.CopyFrom(uuid)

        return service_id_obj

    def GetConfigRule(self, ip_o, ip_d, port_o, port_d):

        config_rule = ConfigRule()

        config_rule_custom = ConfigRule_Custom()

        config_rule.action = ConfigActionEnum.CONFIGACTION_SET

        config_rule_custom.resource_key = "test"

        # config_rule_custom.resource_value = str(self.GenerateRuleValue(ip_o, ip_d, port_o, port_d))

        config_rule_custom.resource_key = "acl"

        config_rule_custom.resource_value = json.dumps(self.GenerateRuleValue(ip_o, ip_d, port_o, port_d))

        config_rule.custom.CopyFrom(config_rule_custom)

        return config_rule

    def SendOutput(self, request, context):

        # SEND CONFIDENCE TO MITIGATION SERVER

        print("Server received mitigation values...", request.confidence, flush=True)

        last_value = request.confidence

        last_tag = request.tag

        LOGGER.info(

            "Attack Mitigator received attack mitigation information. Prediction confidence: %s, Predicted class: %s",

            last_value,

            last_tag,

        )

        ip_o = request.ip_o

        ip_d = request.ip_d

        port_o = request.port_o

        port_d = request.port_d

        # service_id = self.GenerateServiceId(request.service_id)

        # service = GetService(service_id)

        # context_id = self.GenerateContextId("admin")

        sentinel = True

        counter = 0

        # service_id_list = self.ListServiceIds(context_id)

        # print(hello, flush = True)

        # print(hello.service_ids[0].service_uuid.uuid, flush=True)

        # service_id = service_id_list.service_ids[0]

        service_id = request.service_id

        print("Service id: ", service_id, flush=True)

        LOGGER.info("Service id:\n{}".format(service_id))

        while sentinel:

            try:

                sentinel = False

            except Exception as e:

                counter = counter + 1

                print("Waiting 2 seconds", counter, e, flush=True)

                LOGGER.debug("Waiting 2 seconds", counter, e)

                time.sleep(2)

        print("Service obtained from id: ", service, flush=True)

        LOGGER.info("Service obtained from id:\n{}".format(service))

        config_rule = self.GetConfigRule(ip_o, ip_d, port_o, port_d)

        service_config.config_rules.extend([config_rule])

        service.service_config.CopyFrom(service_config)

        print("Service with new rule: ", service, flush=True)

        LOGGER.info("Service with new rule:\n{}".format(service))

        self.UpdateService(service)

        service2 = self.GetService(service_id)

        print("Service obtained from id after updating: ", service2, flush=True)

        LOGGER.info("Service obtained from id after updating:\n{}".format(service2))

        # RETURN OK TO THE CALLER

        return Empty(message=f"OK, received values: {last_tag} with confidence {last_value}.")

    def GetService(self, service_id):

            stub.UpdateService(service)

    def GetMitigation(self, request, context):

        # GET OR PERFORM MITIGATION STRATEGY

        logging.debug("")

        print("Returing mitigation strategy...")

        logging.info("Returning mitigation strategy...")

        k = self.last_value * 2

        return Empty(message=f"Mitigation with double confidence = {k}")