# 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
import logging
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"
SERVICE_CHANNEL = "192.168.165.78:3030"


class l3_attackmitigatorServiceServicerImpl(L3AttackmitigatorServicer):
    def __init__(self):
        LOGGER.debug("Creating Servicer...")
        self.last_value = -1
        self.last_tag = 0

    def GenerateRuleValue(self, ip_o, ip_d, port_o, port_d):
        value = {
            "ipv4:source-address": ip_o,
            "ipv4:destination-address": ip_d,
            "transport:source-port": port_o,
            "transport:destination-port": port_d,
            "forwarding-action": "DROP",
        }
        return value

    def GenerateContextId(self, context_id):
        context_id_obj = ContextId()
        uuid = Uuid()
        uuid.uuid = context_id
        context_id_obj.context_uuid.CopyFrom(uuid)
        return context_id_obj

    def GenerateServiceId(self, service_id):
        service_id_obj = ServiceId()
        context_id = ContextId()
        uuid = Uuid()
        uuid.uuid = 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_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

        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)

        while sentinel:
            try:
                service = self.GetService(service_id)
                sentinel = False
            except Exception as e:
                counter = counter + 1
                print("Waiting 2 seconds", counter, e, flush=True)
                time.sleep(2)

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

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

        service_config = ServiceConfig()
        service_config.config_rules.extend([config_rule])
        service.service_config.CopyFrom(service_config)

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

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

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

    def GetService(self, service_id):
        with grpc.insecure_channel(CONTEXT_CHANNEL) as channel:
            stub = ContextServiceStub(channel)
            return stub.GetService(service_id)

    def ListServiceIds(self, context_id):
        with grpc.insecure_channel(CONTEXT_CHANNEL) as channel:
            stub = ContextServiceStub(channel)
            return stub.ListServiceIds(context_id)

    def UpdateService(self, service):
        with grpc.insecure_channel(SERVICE_CHANNEL) as channel:
            stub = ServiceServiceStub(channel)
            stub.UpdateService(service)

    def GetMitigation(self, request, context):
        # GET OR PERFORM MITIGATION STRATEGY
        logging.debug("")
        print("Returing mitigation strategy...")
        k = self.last_value * 2
        return Empty(message=f"Mitigation with double confidence = {k}")