WatchGuard/label_service.py at main · AsanFillter/WatchGuard · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
import json
import os
from typing import List, Dict, Any
from datetime import datetime


class SimpleLabelManager:

    def __init__(self):
        self.labels_file = "labels.json"
        self._ensure_labels_file()

    def _ensure_labels_file(self):
        if not os.path.exists(self.labels_file):
            default_data = {
                "labels": [],
                "created_at": datetime.now().isoformat(),
                "last_updated": datetime.now().isoformat(),
                "version": "1.0",
            }
            self._save_labels_data(default_data)

    def _load_labels_data(self) -> Dict[str, Any]:
        try:
            with open(self.labels_file, "r", encoding="utf-8") as f:
                data = json.load(f)

                if isinstance(data, list):
                    new_data = {
                        "labels": data,
                        "created_at": datetime.now().isoformat(),
                        "last_updated": datetime.now().isoformat(),
                        "version": "1.0",
                    }
                    self._save_labels_data(new_data)
                    return new_data

                return data

        except (FileNotFoundError, json.JSONDecodeError):
            default_data = {
                "labels": [],
                "created_at": datetime.now().isoformat(),
                "last_updated": datetime.now().isoformat(),
                "version": "1.0",
            }
            self._save_labels_data(default_data)
            return default_data

    def _save_labels_data(self, data: Dict[str, Any]):
        data["last_updated"] = datetime.now().isoformat()
        with open(self.labels_file, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=4, ensure_ascii=False)

    def get_all_labels(self) -> List[str]:
        data = self._load_labels_data()
        labels = data.get("labels", [])

        clean_labels = list(
            set([label.strip() for label in labels if label and label.strip()])
        )
        clean_labels.sort()
        return clean_labels

    def add_label(self, label: str) -> Dict[str, Any]:
        label = label.strip()
        if not label:
            return {"success": False, "message": "Label cannot be empty"}

        if len(label) > 50:
            return {
                "success": False,
                "message": "Label cannot be longer than 50 characters",
            }

        data = self._load_labels_data()
        labels = data.get("labels", [])

        if label in labels:
            return {"success": False, "message": f"Label '{label}' already exists"}

        labels.append(label)
        data["labels"] = labels
        self._save_labels_data(data)

        return {"success": True, "message": f"Label '{label}' added successfully"}

    def remove_label(self, label: str) -> Dict[str, Any]:
        label = label.strip()
        if not label:
            return {"success": False, "message": "Label cannot be empty"}

        data = self._load_labels_data()
        labels = data.get("labels", [])

        if label not in labels:
            return {"success": False, "message": f"Label '{label}' not found"}

        labels.remove(label)
        data["labels"] = labels
        self._save_labels_data(data)

        self._clean_label_from_data(label)

        return {"success": True, "message": f"Label '{label}' removed successfully"}

    def _clean_label_from_data(self, label: str):

        try:
            if os.path.exists("servers.json"):
                with open("servers.json", "r", encoding="utf-8") as f:
                    servers = json.load(f)

                updated = False
                for server_name, server_info in servers.items():
                    if server_info.get("label", "").strip() == label:
                        server_info["label"] = ""
                        updated = True

                if updated:
                    with open("servers.json", "w", encoding="utf-8") as f:
                        json.dump(servers, f, indent=4, ensure_ascii=False)
        except Exception as e:
            print(f"Error cleaning label from servers: {e}")

        try:
            if os.path.exists("domains.json"):
                with open("domains.json", "r", encoding="utf-8") as f:
                    domains = json.load(f)

                updated = False
                for domain_name, domain_info in domains.items():
                    if domain_info.get("label", "").strip() == label:
                        domain_info["label"] = ""
                        updated = True

                if updated:
                    with open("domains.json", "w", encoding="utf-8") as f:
                        json.dump(domains, f, indent=4, ensure_ascii=False)
        except Exception as e:
            print(f"Error cleaning label from domains: {e}")

    def update_dropdown_labels(self):

        return self.get_all_labels()


simple_label_manager = SimpleLabelManager()


def get_labels_for_dropdown() -> List[str]:
    return simple_label_manager.get_all_labels()


def add_new_label(label: str) -> Dict[str, Any]:
    return simple_label_manager.add_label(label)


def remove_existing_label(label: str) -> Dict[str, Any]:
    return simple_label_manager.remove_label(label)