l3_centralizedattackdetectorServiceServicerImpl.py

# Copyright 2021-2023 H2020 TeraFlow (https://www.teraflow-h2020.eu/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function
from datetime import datetime
import os
import grpc
import numpy as np
import onnxruntime as rt
import logging
from common.proto.l3_centralizedattackdetector_pb2 import Empty
from common.proto.l3_centralizedattackdetector_pb2_grpc import L3CentralizedattackdetectorServicer

from common.proto.l3_attackmitigator_pb2 import L3AttackmitigatorOutput
from common.proto.l3_attackmitigator_pb2_grpc import L3AttackmitigatorStub

# KPIs and Monitoring
from common.proto.monitoring_pb2 import KpiDescriptor
from common.proto.kpi_sample_types_pb2 import KpiSampleType

# from monitoring.client.MonitoringClient import MonitoringClient
from monitoring.client.MonitoringClient import MonitoringClient
from common.proto.monitoring_pb2 import Kpi

from common.tools.timestamp.Converters import timestamp_utcnow_to_float
from common.proto.context_pb2 import Timestamp

LOGGER = logging.getLogger(__name__)
here = os.path.dirname(os.path.abspath(__file__))
MODEL_FILE = os.path.join(here, "ml_model/crypto_5g_rf_spider_features.onnx")

classification_threshold = os.getenv("CAD_CLASSIFICATION_THRESHOLD", 0.5)


class l3_centralizedattackdetectorServiceServicerImpl(L3CentralizedattackdetectorServicer):
    def __init__(self):
        LOGGER.debug("Creating Servicer...")
        self.inference_values = []
        self.model = rt.InferenceSession(MODEL_FILE)
        self.input_name = self.model.get_inputs()[0].name
        self.label_name = self.model.get_outputs()[0].name
        self.prob_name = self.model.get_outputs()[1].name
        self.monitoring_client = MonitoringClient()
        self.predicted_class_kpi_id = None
        self.class_probability_kpi_id = None

    def create_predicted_class_kpi(self, client: MonitoringClient, service_id):
        # create kpi
        kpi_description: KpiDescriptor = KpiDescriptor()
        kpi_description.kpi_description = "L3 security status of service {}".format(service_id)
        # kpi_description.service_id.service_uuid.uuid = service_id
        kpi_description.service_id.service_uuid.uuid = str(service_id)
        kpi_description.kpi_sample_type = KpiSampleType.KPISAMPLETYPE_UNKNOWN
        new_kpi = client.SetKpi(kpi_description)

        LOGGER.info("Created Predicted Class KPI {}...".format(new_kpi.kpi_id))

        return new_kpi

    def create_class_prob_kpi(self, client: MonitoringClient, service_id):
        # create kpi
        kpi_description: KpiDescriptor = KpiDescriptor()
        kpi_description.kpi_description = "L3 security status of service {}".format(service_id)
        kpi_description.service_id.service_uuid.uuid = str(service_id)
        kpi_description.kpi_sample_type = KpiSampleType.KPISAMPLETYPE_UNKNOWN
        new_kpi = client.SetKpi(kpi_description)

        LOGGER.info("Created Class Probability KPI {}...".format(new_kpi.kpi_id))

        return new_kpi

    def make_inference(self, request):
        # ML MODEL
        x_data = np.array(
            [
                [
                    request.c_pkts_all,
                    request.c_ack_cnt,
                    request.c_bytes_uniq,
                    request.c_pkts_data,
                    request.c_bytes_all,
                    request.s_pkts_all,
                    request.s_ack_cnt,
                    request.s_bytes_uniq,
                    request.s_pkts_data,
                    request.s_bytes_all,
                ]
            ]
        )

        predictions = self.model.run([self.prob_name], {self.input_name: x_data.astype(np.float32)})[0]
        # Output format
        output_message = {
            "confidence": None,
            "timestamp": datetime.now().strftime("%d/%m/%Y %H:%M:%S"),
            "ip_o": request.ip_o,
            "ip_d": request.ip_d,
            "tag_name": None,
            "tag": None,
            "flow_id": request.flow_id,
            "protocol": request.protocol,
            "port_o": request.port_o,
            "port_d": request.port_d,
            "ml_id": "RandomForest",
            "service_id": request.service_id,
            "time_start": request.time_start,
            "time_end": request.time_end,
        }
        if predictions[0][1] >= classification_threshold:
            output_message["confidence"] = predictions[0][1]
            output_message["tag_name"] = "Crypto"
            output_message["tag"] = 1
        else:
            output_message["confidence"] = predictions[0][0]
            output_message["tag_name"] = "Normal"
            output_message["tag"] = 0

        return L3AttackmitigatorOutput(**output_message)

    def SendInput(self, request, context):
        # PERFORM INFERENCE WITH SENT INPUTS
        logging.debug("")
        print("Inferencing ...", flush=True)

        # STORE VALUES
        self.inference_values.append(request)

        # MAKE INFERENCE
        output = self.make_inference(request)

        # Monitoring
        service_id = request.service_id

        if self.predicted_class_kpi_id is None:
            self.predicted_class_kpi_id = self.create_predicted_class_kpi(self.monitoring_client, service_id)

        if self.class_probability_kpi_id is None:
            self.class_probability_kpi_id = self.create_class_prob_kpi(self.monitoring_client, service_id)

        # Packet -> DAD -> CAD -> ML -> (2 Instantaneous Value: higher class probability, predicted class) -> Monitoring
        # In addition, two counters:
        # Counter 1: Total number of crypto attack connections
        # Counter 2: Rate of crypto attack connections with respect to the total number of connections

        kpi_class = Kpi()
        kpi_class.kpi_id.kpi_id.uuid = str(self.predicted_class_kpi_id)
        kpi_class.kpi_value.int32Val = 1 if output.tag_name == "Crypto" else 0

        kpi_prob = Kpi()
        kpi_prob.kpi_id.kpi_id.uuid = str(self.class_probability_kpi_id)
        kpi_prob.kpi_value.floatVal = output.confidence

        # timestamp = timestamp_utcnow_to_float()
        timestamp = Timestamp()
        timestamp.timestamp = timestamp_utcnow_to_float()

        kpi_class.timestamp.CopyFrom(timestamp)
        kpi_prob.timestamp.CopyFrom(kpi_class.timestamp)

        self.monitoring_client.IncludeKpi(kpi_class)
        self.monitoring_client.IncludeKpi(kpi_prob)

        if output.tag_name == "Crypto":
            # SEND INFO TO MITIGATION SERVER
            try:
                with grpc.insecure_channel("192.168.165.78:10002") as channel:
                    stub = L3AttackmitigatorStub(channel)
                    print("Sending to mitigator...", flush=True)
                    response = stub.SendOutput(output)
                    # print("Response received", response, "Hola", flush=True)
                    # print("Sent output to mitigator and received: ", response.message) #FIX No message received

                    # RETURN "OK" TO THE CALLER
                return Empty(message="OK, information received and mitigator notified abou the attack")
            except Exception as e:
                print("This is an exception", repr(e), flush=True)
                print("Couldnt find l3_attackmitigator", flush=True)
                return Empty(message="Mitigator Not found")
        else:
            print("No attack detected", flush=True)
            return Empty(message="OK, information received (no attack detected)")

    def GetOutput(self, request, context):
        logging.debug("")
        print("Returing inference output...")
        k = np.multiply(self.inference_values, [2])
        k = np.sum(k)
        return self.make_inference(k)