"""
Dieses Skript sammelt Informationen über alle Docker-Container auf dem Host.
Es exportiert die Container- und Netzwerkdaten als CSV-Datei und erzeugt eine GraphML-Datei,
die die Container und Netzwerke als Knoten sowie deren Verbindungen als Kanten darstellt.
Die GraphML-Datei ist für die Visualisierung in yEd optimiert (Star-Schema).
"""

import docker
import pandas as pd
import re
import os
import datetime
import socket

def sanitize_id(text):
    """
    Erzeugt eine gültige ID für GraphML-Knoten.
    Entfernt alle Zeichen außer Buchstaben, Zahlen und Unterstrichen.
    IDs dürfen nicht mit einer Zahl beginnen.
    """
    return re.sub(r'\W|^(?=\d)', '_', text)

def get_container_info():
    """
    Sammelt Informationen über alle Docker-Container.
    Gibt eine Liste von Dictionaries zurück, die folgende Felder enthalten:
    - Container Name
    - Watchtower Enabled (aus den Labels com.centurylinklabs.watchtower.enable oder watchtower.enable)
    - Network (Name des Netzwerks)
    - IP Address (IP-Adresse im Netzwerk)
    - Gateway (Gateway-Adresse im Netzwerk)
    """
    client = docker.from_env()
    containers = client.containers.list(all=True)

    data = []

    for container in containers:
        container_info = client.api.inspect_container(container.id)

        name = container.name
        labels = container_info.get('Config', {}).get('Labels', {})
        watchtower_enabled = labels.get('com.centurylinklabs.watchtower.enable',
                                        labels.get('watchtower.enable', 'not set'))

        networks = container_info.get('NetworkSettings', {}).get('Networks', {})

        for net_name, net_data in networks.items():
            ip_address = net_data.get('IPAddress', 'N/A')
            gateway = net_data.get('Gateway', 'N/A')

            data.append({
                'Container Name': name,
                'Watchtower Enabled': watchtower_enabled,
                'Network': net_name,
                'IP Address': ip_address,
                'Gateway': gateway
            })

    return data

def save_to_csv(data, filename='docker_container_info.csv'):
    """
    Speichert die gesammelten Container-Informationen als CSV-Datei.
    :param data: Liste von Dictionaries mit Container-Infos
    :param filename: Name der zu erzeugenden CSV-Datei
    """
    df = pd.DataFrame(data)
    df.to_csv(filename, index=False)
    print(f"[+] Daten erfolgreich in {filename} gespeichert.")

def save_to_graphml(data, filename='docker_container_info.graphml'):
    """
    Erstellt eine GraphML-Datei für yEd, die die Container und Netzwerke als Knoten und deren Verbindungen als Kanten darstellt.
    - Netzwerke werden zentral angeordnet (Star Schema)
    - Container werden kreisförmig um die Netzwerke positioniert
    - Die Labels der Container enthalten Name, Watchtower-Status und IP-Infos
    :param data: Liste von Dictionaries mit Container-Infos
    :param filename: Name der zu erzeugenden GraphML-Datei
    """
    from xml.etree.ElementTree import Element, SubElement, ElementTree, register_namespace

    # Nur yEd Namespace registrieren!
    register_namespace('y', "http://www.yworks.com/xml/graphml")

    graphml = Element('graphml', xmlns="http://graphml.graphdrawing.org/xmlns")

    # yEd-Label-Schlüssel
    key_label = SubElement(graphml, 'key', id="d0", **{
        "for": "node",
        "yfiles.type": "nodegraphics"
    })
    key_watchtower = SubElement(graphml, 'key', id="d1", **{
        "for": "node",
        "attr.name": "Watchtower Enabled",
        "attr.type": "string"
    })

    graph = SubElement(graphml, 'graph', id="G", edgedefault="undirected")

    container_labels = {}
    containers = set()
    networks = set()
    for entry in data:
        cname = entry['Container Name']
        containers.add(cname)
        networks.add(entry['Network'])
        if cname not in container_labels:
            container_labels[cname] = {}
        container_labels[cname]['Watchtower Enabled'] = entry['Watchtower Enabled']

    # Star Schema: Netzwerke in der Mitte, Container im Kreis
    import math

    # Positioniere Netzwerk-Knoten im Zentrum
    center_x = 500
    center_y = 400
    net_radius = 0  # Alle Netzwerke exakt im Zentrum
    net_positions = {}
    for idx, n in enumerate(networks):
        node_id = f"network_{sanitize_id(n)}"
        node = SubElement(graph, 'node', id=node_id)
        net_infos = [entry for entry in data if entry['Network'] == n]
        net_texts = [f"{entry['Container Name']}: IP: {entry['IP Address']}, GW: {entry['Gateway']}" for entry in net_infos]
        label_text = f"{n}\n" + "\n".join(net_texts)
        data_elem = SubElement(node, 'data', key="d0")
        y_shape_node = SubElement(data_elem, '{http://www.yworks.com/xml/graphml}ShapeNode')
        # Alle Netzwerke im Zentrum
        geometry = SubElement(y_shape_node, '{http://www.yworks.com/xml/graphml}Geometry',
                              x=str(center_x), y=str(center_y), width="200", height="60")
        y_node_label = SubElement(y_shape_node, '{http://www.yworks.com/xml/graphml}NodeLabel')
        y_node_label.text = label_text
        net_positions[n] = (center_x, center_y)

    # Container-Knoten im Kreis um das Zentrum
    container_count = len(containers)
    radius = 300
    angle_step = 2 * math.pi / max(container_count, 1)
    container_positions = {}
    for idx, c in enumerate(containers):
        node_id = f"container_{sanitize_id(c)}"
        ip_infos = [entry for entry in data if entry['Container Name'] == c]
        ip_texts = [f"IP: {info['IP Address']}, GW: {info['Gateway']}, Net: {info['Network']}" for info in ip_infos]
        # Ermittle Watchtower-Status aus com.centurylinklabs.watchtower.enable
        watchtower_status = "aktiviert" if any(
            info.get('Watchtower Enabled', '').lower() in ['true', '1', 'yes', 'enabled']
            for info in ip_infos
        ) else "deaktiviert"
        label_text = f"{c}\nWatchtower: {watchtower_status}\n" + "\n".join(ip_texts)
        node = SubElement(graph, 'node', id=node_id)
        data_elem = SubElement(node, 'data', key="d0")
        y_shape_node = SubElement(data_elem, '{http://www.yworks.com/xml/graphml}ShapeNode')
        # Kreisförmige Position
        angle = idx * angle_step
        x = center_x + radius * math.cos(angle)
        y = center_y + radius * math.sin(angle)
        geometry = SubElement(y_shape_node, '{http://www.yworks.com/xml/graphml}Geometry',
                              x=str(x), y=str(y), width="200", height="60")
        y_node_label = SubElement(y_shape_node, '{http://www.yworks.com/xml/graphml}NodeLabel')
        y_node_label.text = label_text
        watchtower_elem = SubElement(node, 'data', key="d1")
        watchtower_elem.text = container_labels[c]['Watchtower Enabled']
        container_positions[c] = (x, y)

    edge_id = 0
    for entry in data:
        source = f"container_{sanitize_id(entry['Container Name'])}"
        target = f"network_{sanitize_id(entry['Network'])}"
        edge = SubElement(graph, 'edge', id=f"e{edge_id}", source=source, target=target)
        edge_id += 1

    tree = ElementTree(graphml)
    tree.write(filename, encoding='utf-8', xml_declaration=True)
    print(f"[+] GraphML erfolgreich in {filename} gespeichert.")

if __name__ == '__main__':
    print("[*] Sammle Container-Informationen...")
    container_data = get_container_info()

    # Erzeuge standardisierte Dateinamen: containers_snapshot_<YYYY-MM-DD>_<hostname>.(csv|graphml)
    date_str = datetime.date.today().isoformat()
    hostname = sanitize_id(socket.gethostname())
    base_name = f"containers_snapshot_{date_str}_{hostname}"

    out_dir = 'out_01'
    try:
        os.makedirs(out_dir, exist_ok=True)
    except Exception:
        # Falls das Anlegen fehlschlägt, fahren wir mit dem aktuellen Verzeichnis fort
        out_dir = '.'

    csv_file = os.path.join(out_dir, f"{base_name}.csv")
    graph_file = os.path.join(out_dir, f"{base_name}.graphml")

    save_to_csv(container_data, filename=csv_file)
    save_to_graphml(container_data, filename=graph_file)
    print(f"[+] Dateien: {csv_file}, {graph_file}")