import subprocess
import yaml
import os
import ruamel.yaml
import ipaddress


def execute_command (command):
    try:
        subprocess.check_call(command, shell=True)
        print("Command executed")
    except subprocess.CalledProcessError as e:
        print ("Error", e)

def createDeploymentFolder(name):
    """
    Creates the new folder where place all the yaml files to use.

    Returns:
        Path Directory just created.
    """
    result = name.split("outputs_")[-1]
    newDirectoryPath = os.path.join(os.getcwd(), "nse-composition-" +result)
   
    if not os.path.exists(newDirectoryPath):
        os.makedirs(newDirectoryPath)
    return newDirectoryPath

def getFolderName():
    """
    Retrieves folder name where place all the yaml files.
    
    Returns:
        FolderName
    """
    script_directory = os.path.dirname(os.path.abspath(__file__))
    api_directory = os.path.join(script_directory, "API")

    latest_setup = getLastSetup(api_directory)
    new_directory_path = createDeploymentFolder(latest_setup)
    return new_directory_path

def getLastSetup(directory_path):
    """
    Retrieves the ID(date + Random_String) from the security intent received by the API.
    In this way it is possible create and get the folder name used then to deploy the cluster.
    
    Returns:
        FolderName
    """
    subdirectories = [d for d in os.listdir(directory_path) if os.path.isdir(os.path.join(directory_path, d))]

    if not subdirectories:
        return None 

    subdirectories.sort(key=lambda d: os.path.getctime(os.path.join(directory_path, d)), reverse=True)
    latest_directory = os.path.join(directory_path, subdirectories[0])
    return os.path.basename(latest_directory)
    
def generateNamespace ():
    folderName = getFolderName()
    yaml = ruamel.yaml.YAML()
    namespace_data = {
        "apiVersion": "v1",
        "kind": "Namespace",
        "metadata": {
            "name": "ns-nse-composition"
        }
    }
    filename="namespace.yaml"
    with open(os.path.join(folderName+"/"+filename), "w") as file:
        yaml.indent(offset=2) 
        yaml.dump(namespace_data, file)
    configureKustomization(filename)

def generateNetworkService():
    """
    Generates a NetworkService YAML definition based on input data and writes it to a file.
    
    This function reads input data from an input YAML file, processes it, and creates a NetworkService
    YAML definition based on that data. It then writes the generated NetworkService definition to an output
    YAML file and potentially calls a separate function to configure the kustomization file.

    Returns:
        None
    """
    script_directory = os.path.dirname(os.path.abspath(__file__))
    api_directory = os.path.join(script_directory, "API")
    latest_setup = getLastSetup(api_directory)
    obj_yaml_path = os.path.join(api_directory, latest_setup, "obj.yaml")

    yaml = ruamel.yaml.YAML()
    #This is needed to format correctly the file 
    yaml.indent(mapping=2, sequence=4, offset=2)
    with open(obj_yaml_path, 'r') as file:
        input_data = yaml.load(file)
    #First part of the file
    output_data = {
        'apiVersion': 'networkservicemesh.io/v1',
        'kind': 'NetworkService',
        'metadata': {
            # this is the name of the network service. Change
            # security-chain into something else if you want a different name
            'name': 'security-chain' 
        },
        'spec': {
            'payload': 'ETHERNET',
            'matches': []
        }
    }
    
    app_values = [security_function['name'] for security_function in input_data['spec']['securityFunctions']]

    for app in app_values[:1]:
        output_data['spec']['matches'].append({
            'routes': [{
                'destination_selector': {
                    'app': app
                }
            }]
        })
    #specific section for the routes with source and destination selector
    for i, app in enumerate(app_values):
        match = {
            'source_selector': {
                     'app': app
            },
                'routes': [

            ]
        }
        #Checks if it is the last of the list
        if i < len(app_values) - 1:
            match['routes'].append({
                'destination_selector': {
                         'app': app_values[i+1]
                }
            })
        else:
            #configure the last routes to the final endpoint
            match['routes'].append({
                'destination_selector': {
                         'app': 'gateway'
                }
            })
        output_data['spec']['matches'].append(match)

    #configure the first route, the default one, taking the first element of the list of the security functions


    filename = "ns.yaml"
    folderName = getFolderName()
    with open(folderName+"/ns.yaml", 'w') as file:
        yaml.dump(output_data, file)
    
    configureKustomization(filename)  

def generateTemplate():
    """
    Generates Kubernetes Deployment definitions based on input data and configures them.

    This code block reads input data from an input YAML file, processes it to generate Kubernetes Deployment definitions
    for each security function, and writes the deployment templates to separate files. It also calls the `addContainer()`
    function for each deployment to configure and add container definitions to the templates.

    Returns:
        None
    """
    script_directory = os.path.dirname(os.path.abspath(__file__))
    api_directory = os.path.join(script_directory, "API")
    latest_setup = getLastSetup(api_directory)
    obj_yaml_path = os.path.join(api_directory, latest_setup, "obj.yaml")

    with open(obj_yaml_path, 'r') as file:
        input_data = ruamel.yaml.YAML().load(file)
    
    serviceName = getServiceName()
    
    # Define a base IPv4 subnet and a subnet generator for assigning subnets to security functions
    base_subnet = ipaddress.IPv4Network('172.16.0.0/16')
    subnet_generator = base_subnet.subnets(new_prefix=31)

    #Here is where additional fields should be handled, as well as adding the related fields in the deployment.    
    for security_function in input_data.get('spec', {}).get('securityFunctions', []):
        name = security_function.get('name')
        if name:
            name = name.lower()
        image = security_function.get('image') 
        subnet = str(next(subnet_generator))
        
        #should be handled 
        port = security_function.get('port')

        #This is the standard deployment that the endpoint should have in order to handle the 2 Kernel Interfaces
        #More information about the NSM details are provided in the Readme
        deployment_template = {
            "apiVersion": "apps/v1",
            "kind": "Deployment",
            "metadata": {
                "name": f"{name}-endpoint",
                "labels": {
                    "app": name
                }
            },
            "spec": {
                "selector": {
                    "matchLabels": {
                        "app": name
                    }
                },
                "template": {
                    "metadata": {
                        "labels": {
                            "app": name,
                            "spiffe.io/spiffe-id": "true"
                        }
                    },
                    "spec": {
                        "containers": [
                            {
                                "name": "nse",
                                "image": "ghcr.io/networkservicemesh/ci/cmd-nse-icmp-responder:6235cf9",
                                "imagePullPolicy": "IfNotPresent",
                                "securityContext": {
                                    "privileged": True
                                },
                                "env": [
                                    {
                                        "name": "SPIFFE_ENDPOINT_SOCKET",
                                        "value": "unix:///run/spire/sockets/agent.sock"
                                    },
                                    {
                                        "name": "NSM_LABELS",
                                        "value": f"app:{name}"
                                    },
                                    {
                                        "name": "NSM_SERVICE_NAMES",
                                        "value": serviceName,
                                    },
                                    {
                                        "name": "NSM_CIDR_PREFIX",
                                        "value": subnet
                                    },
                                    {
                                        "name": "NSM_NAME",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "metadata.name"
                                            }
                                        }
                                    },
                                    {
                                        "name": "NSM_CONNECT_TO",
                                        "value": "unix:///var/lib/networkservicemesh/nsm.io.sock"
                                    },
                                    {
                                        "name": "NSM_REGISTER_SERVICE",
                                        "value": "false"
                                    }
                                ],
                                "volumeMounts": [
                                    {
                                        "name": "spire-agent-socket",
                                        "mountPath": "/run/spire/sockets",
                                        "readOnly": True
                                    },
                                    {
                                        "name": "nsm-socket",
                                        "mountPath": "/var/lib/networkservicemesh",
                                        "readOnly": True
                                    }
                                ],
                                "resources": {
                                    "requests": {
                                        "cpu": "100m",
                                        "memory": "40Mi"
                                    },
                                    "limits": {
                                        "memory": "80Mi",
                                        "cpu": "200m"
                                    }
                                }
                            },
                            {
                                "name": "cmd-nsc",
                                "image": "ghcr.io/networkservicemesh/ci/cmd-nsc:d105640",
                                "imagePullPolicy": "IfNotPresent",
                                "env": [
                                    {
                                        "name": "SPIFFE_ENDPOINT_SOCKET",
                                        "value": "unix:///run/spire/sockets/agent.sock"
                                    },
                                    {
                                        "name": "NSM_REGISTER_SERVICE",
                                        "value": "false"
                                    },
                                    {
                                        "name": "NSM_NETWORK_SERVICES",
                                        "value": f"kernel://security-chain/kernelint?app={name}"
                                    },
                                    {
                                        "name": "NSM_LOG_LEVEL",
                                        "value": "TRACE"
                                    },
                                    {
                                        "name": "NSM_CONNECT_TO",
                                        "value": "unix:///var/lib/networkservicemesh/nsm.io.sock"
                                    }
                                ],
                                "volumeMounts": [
                                    {
                                        "name": "spire-agent-socket",
                                        "mountPath": "/run/spire/sockets",
                                        "readOnly": True
                                    },
                                    {
                                        "name": "nsm-socket",
                                        "mountPath": "/var/lib/networkservicemesh",
                                        "readOnly": True
                                    }
                                ],
                                "resources": {
                                    "requests": {
                                        "cpu": "100m",
                                        "memory": "40Mi"
                                    },
                                    "limits": {
                                        "memory": "80Mi",
                                        "cpu": "200m"
                                      }
                                }
                            }
                        ],
                
                        "volumes": [
                            {
                                "name": "spire-agent-socket",
                                "hostPath": {
                                    "path": "/run/spire/sockets",
                                    "type": "Directory"
                                }
                            },
                            {
                                "name": "nsm-socket",
                                "hostPath": {
                                    "path": "/var/lib/networkservicemesh",
                                    "type": "DirectoryOrCreate"
                                }
                            }
                    ]
                },
                },
            }
        }
        folderName = getFolderName()
        with open(folderName+f"/deployment_template_{name}.yaml", 'w') as file:
                  ruamel.yaml.YAML().dump(deployment_template, file)
        addContainer(name, image)

def configureKustomization(filename):
    """
    Configures a Kustomization file for managing Kubernetes resources and adds a resource entry.

    This function is responsible for managing a Kustomization file, which is used to manage
    and customize Kubernetes resources. It checks if the Kustomization file already exists.
    If it exists, it loads and updates it; otherwise, it creates a new one. It adds the provided
    filename as a resource entry in the Kustomization file.

    Args:
        filename (str): The filename to be added as a resource in the Kustomization file.

    Returns:
        None
    """
    folderName = getFolderName()
    if os.path.exists(folderName+"/kustomization.yaml"):
        with open(folderName+"/kustomization.yaml", "r") as kustomization_file:
            kustomization = ruamel.yaml.YAML().load(kustomization_file)
    else:
        kustomization = {
            "apiVersion": "kustomize.config.k8s.io/v1beta1",
            "kind": "Kustomization",
            "namespace": "ns-nse-composition",
            "resources": [],
            }
    if filename not in kustomization["resources"]:
        kustomization["resources"].append(filename)

    with open(folderName+"/kustomization.yaml", "w") as kustomization_file:
        ruamel.yaml.YAML().dump(kustomization, kustomization_file)

def addContainer(name, image):
    """
    Adds a container definition to a Kubernetes Deployment and writes it to a file.

    This function opens a Kubernetes Deployment template file, adds a new container definition
    to the template, and then writes the updated template to a new file. It calls `configureKustomization()` 
    function to write the filename in the kustomization file.
    If the image is "Squid" Then there is an additional configuration to use it.

    Args:
        name (str): The name of the container.
        image (str): The Docker image for the container.

    Returns:
        None
    """
    folderName = getFolderName()
    with open(folderName+f"/deployment_template_{name}.yaml", 'r') as file:
        deployment_template = ruamel.yaml.YAML().load(file)
    #If it is needed, the changes to configure the 3° container should be applied here.
    #Like adding a new param or specification
    container = {
        'name': name,
        'image': image,
        'imagePullPolicy': 'IfNotPresent',
        'securityContext':{
            'privileged': True},
        'command': ["sh", "-c", "sleep infinity"]
    }
    
        
    ################################################
    # This part is specific for the suricata image #
    ################################################
    if image == "jasonish/suricata":
        container['volumeMounts'] = [
            {
                'name': 'suricata-rules',
                'mountPath': '/additional/rules/local',
                'readOnly': False,
          
            }]
        volume = [{
     'name': 'suricata-rules',
    'configMap': {
        'name': 'suricata-rules'
                }
        }]
        deployment_template['spec']['template']['spec']['volumes'].extend(volume)
        configureSuricata()


    #############################################
    # This part is specific for the squid image #
    #############################################
    if image == "ubuntu/squid:edge":
        container['ports'] = [
            {
                'containerPort': 3128,
                'name': 'squid',
                'protocol': 'TCP'
            }
        ]
        container['volumeMounts'] = [
            {
                'name': 'squid-config-volume',
                'mountPath': '/etc/squid/squid.conf',
                'subPath': 'squid.conf'
            },
            {
                'name': 'squid-data',
                'mountPath': '/var/spool/squid'
            }
        ]

        volume = [{
     'name': 'squid-config-volume',
    'configMap': {
        'name': 'squid-config'
                }
        },
        {
    'name': 'squid-data',
    'persistentVolumeClaim': {
        'claimName': 'squid-volume-claim'
            }
        }]
        
        deployment_template['spec']['template']['spec']['volumes'].extend(volume)

        configureSquid()
      

    #Appends the container description in the specific section of the yaml file.
    deployment_template['spec']['template']['spec']['containers'].append(container)
    filename = f"{name}.yaml"

    with open(folderName+f"/{name}.yaml", 'w') as file:
        ruamel.yaml.YAML().dump(deployment_template, file)
    configureKustomization(filename)

def getServiceName():
    #Retrieve the service name
    folderName = getFolderName()
    with open(folderName+"/ns.yaml", 'r') as file:
        yaml_data = yaml.load(file, Loader=yaml.FullLoader)
    
    if 'metadata' in yaml_data and 'name' in yaml_data['metadata']:
        return yaml_data['metadata']['name']
    else:
        return None

#############################################################
# All the following functions are used to configure the squid container


def configureSquid():
    createStorageSquid()
    create_PVC_Squid()
    createConfFile()
    patchKustomizationFileSquid()
    return

def createStorageSquid():
    """
    Creates the storage yaml file needed to configure squid. 
    At the end, adds the filename to the kustomization file.

    Returns:
        None
    """
    folderName=getFolderName()
    yaml_data = [
        {
            "apiVersion": "storage.k8s.io/v1",
            "kind": "StorageClass",
            "metadata": {
                "name": "kind-hostpath"
            },
            "provisioner": "docker.io/hostpath"
        },
        {
            "apiVersion": "v1",
            "kind": "PersistentVolume",
            "metadata": {
                "name": "kind-storage"
            },
            "spec": {
                "capacity": {
                    "storage": "600Mi"
                },
                "accessModes": [
                    "ReadWriteOnce"
                ],
                "persistentVolumeReclaimPolicy": "Retain",
                "storageClassName": "kind-hostpath",
                "hostPath": {
                    "path": "/mnt"
                }
            }
        }
    ]
    filename="storage.yaml"
    with open(folderName+"/"+filename, "w") as yaml_file:
        yaml.dump_all(yaml_data, yaml_file, default_flow_style=False)
    configureKustomization(filename)

def create_PVC_Squid():
    folderName=getFolderName()
    yaml_data = {
        "apiVersion": "v1",
        "kind": "PersistentVolumeClaim",
        "metadata": {
            "name": "squid-volume-claim"
        },
        "spec": {
            "accessModes": [
                "ReadWriteOnce"
            ],
            "resources": {
                "requests": {
                    "storage": "500Mi"
                }
            },
            "storageClassName": "kind-hostpath",
            "volumeName": "kind-storage"
        }
    }
    filename="pvc.yaml"
    with open(folderName+"/"+filename, "w") as yaml_file:
        yaml.dump(yaml_data, yaml_file, default_flow_style=False)   
    configureKustomization(filename)     

def createConfFile():
    """
    Creates the squid.conf file. This is the standard file provided by Squid. 
    Check https://hub.docker.com/r/ubuntu/squid to see how perform changes.    
    Returns:
        None
    """
    folderName=getFolderName()
    config_content = """\
acl localnet src 0.0.0.1-0.255.255.255	# RFC 1122 "this" network (LAN)
acl localnet src 10.0.0.0/8		# RFC 1918 local private network (LAN)
acl localnet src 100.64.0.0/10		# RFC 6598 shared address space (CGN)
acl localnet src 169.254.0.0/16 	# RFC 3927 link-local (directly plugged) machines
acl localnet src 172.16.0.0/12		# RFC 1918 local private network (LAN)
acl localnet src 192.168.0.0/16		# RFC 1918 local private network (LAN)
acl localnet src fc00::/7       	# RFC 4193 local private network range
acl localnet src fe80::/10      	# RFC 4291 link-local (directly plugged) machines
acl SSL_ports port 443
acl Safe_ports port 80		# http
acl Safe_ports port 21		# ftp
acl Safe_ports port 443		# https
acl Safe_ports port 70		# gopher
acl Safe_ports port 210		# wais
acl Safe_ports port 1025-65535	# unregistered ports
acl Safe_ports port 280		# http-mgmt
acl Safe_ports port 488		# gss-http
acl Safe_ports port 591		# filemaker
acl Safe_ports port 777		# multiling http
acl special_ips src  172.16.0.0/12
acl CONNECT method CONNECT
http_access deny special_ips
http_access deny !Safe_ports
http_access deny CONNECT !SSL_ports
http_access allow localhost manager
http_access deny manager
http_access allow localhost
http_access allow localnet
http_access deny all
http_port 3128
coredump_dir /var/spool/squid
refresh_pattern ^ftp:		1440	20%	10080
refresh_pattern ^gopher:	1440	0%	1440
refresh_pattern -i (/cgi-bin/|\?) 0	0%	0
refresh_pattern \/(Packages|Sources)(|\.bz2|\.gz|\.xz)$ 0 0% 0 refresh-ims
refresh_pattern \/Release(|\.gpg)$ 0 0% 0 refresh-ims
refresh_pattern \/InRelease$ 0 0% 0 refresh-ims
refresh_pattern \/(Translation-.*)(|\.bz2|\.gz|\.xz)$ 0 0% 0 refresh-ims
refresh_pattern .		0	20%	4320
logfile_rotate 0
"""
    filename="squid.conf"
    with open(folderName+"/"+filename, "w") as config_file:
        config_file.write(config_content)


def patchKustomizationFileSquid():
    """
    Patch the kustomization file adding the part related to the config map: in this 
    way squid.conf can be applied.  
    Returns:
        None
    """
    
    folderName = getFolderName()
    yaml = ruamel.yaml.YAML()
    ruamel.yaml.YAML().indent(offset=2, mapping=2, sequence=2)
    yaml.preserve_quotes = True
    with open(folderName+"/kustomization.yaml", 'r') as file:
        kustomizationFile = yaml.load(file)
    if 'configMapGenerator' not in kustomizationFile:
            kustomizationFile['configMapGenerator'] = [    
            {"name": "squid-config",
            "files": ["squid.conf"]}]
                 
    configMap = {"name": "squid-config",
        "files": ["squid.conf"]}
                 
    if not any(config == configMap for config in kustomizationFile['configMapGenerator']):
        kustomizationFile['configMapGenerator'].append(configMap)
 
    with open(folderName+"/kustomization.yaml", 'w') as file:
        yaml.dump(kustomizationFile, file)

#####################################################
# This section is related to suricata configuration #             
#####################################################
def configureSuricata():
    createRulesFile()
    patchKustomizationFileSuricata()
    return

def createRulesFile():
    """
    This is the function used to create the file where the rules for suricata are placed.
    Add here the new rules.
    """
    folderName=getFolderName()
    config_content = """\
        alert icmp $HOME_NET any -> $HOME_NET any (msg:"Rilevato traffico ICMP pericoloso"; sid:1;)
        alert icmp any any -> any any (msg:"Altra prova di regole suricata"; sid:2;)
        """
    filename="suricata.rules"
    with open(folderName+"/"+filename, "w") as config_file:
        config_file.write(config_content)
    
def patchKustomizationFile(filename):
    """
    Update the 'kustomization.yaml' configuration file by adding a filename as a patch.

    Args:
    - filename (str): The name of the file to be added as a patch to 'kustomization.yaml'.

    This function modifies the 'kustomization.yaml' file located in a specific folder. It adds the given filename as a patch to the 'patchesStrategicMerge' list within the 'kustomization.yaml' file if it's not already present.

    """
    folderName = getFolderName()
    yaml = ruamel.yaml.YAML()
    ruamel.yaml.YAML().indent(offset=2, mapping=2, sequence=2)
    yaml.preserve_quotes = True
    with open(folderName+"/kustomization.yaml", 'r') as file:
        kustomizationFile = yaml.load(file)
    if 'patchesStrategicMerge' not in kustomizationFile:
        kustomizationFile['patchesStrategicMerge'] = [filename]
    elif filename not in kustomizationFile['patchesStrategicMerge']:
        kustomizationFile['patchesStrategicMerge'].append(filename)

    with open(folderName+"/kustomization.yaml", 'w') as file:
        yaml.dump(kustomizationFile, file)


def patchKustomizationFileSuricata():
    """
    Patch the kustomization file adding the part related to the config map: in this 
    way surica.rules can be applied.  
    Returns:
        None
    """
    folderName = getFolderName()
    yaml = ruamel.yaml.YAML()
    ruamel.yaml.YAML().indent(offset=2, mapping=2, sequence=2)
    yaml.preserve_quotes = True
    with open(folderName+"/kustomization.yaml", 'r') as file:
        kustomizationFile = yaml.load(file)
    if 'configMapGenerator' not in kustomizationFile:
            kustomizationFile['configMapGenerator'] = [    
            {"name": "suricata-rules",
            "files": ["suricata.rules"]}]
                 
    configMap = {"name": "suricata-rules",
        "files": ["suricata.rules"]}
                 
    if not any(config == configMap for config in kustomizationFile['configMapGenerator']):
        kustomizationFile['configMapGenerator'].append(configMap)
 
    with open(folderName+"/kustomization.yaml", 'w') as file:
        yaml.dump(kustomizationFile, file)

