Enterprise Policy Manager GUI.py

# -*- coding: utf-8 -*-
"""
Fantastic v5.7 – Enterprise Policy Manager + Super Perfect Simulator
PySide6 GUI with Redis backend, full simulator integrated with user dropdown.
Supports staged changes and highlights for modified/deleted items.
"""

import sys, json, uuid, redis, ipaddress
from urllib.parse import urlparse
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QListWidget, QListWidgetItem,
    QPushButton, QLabel, QLineEdit, QTextEdit, QCheckBox,
    QVBoxLayout, QHBoxLayout, QSplitter, QDialog, QFormLayout,
    QDialogButtonBox, QMessageBox, QInputDialog, QFileDialog, QComboBox,
    QProgressDialog, QMenu
)
from PySide6.QtCore import Qt, QThread, Signal, QTimer, QEvent
from PySide6.QtGui import QColor, QPalette


# ---------------- JsonWorker Class -----------------
class JsonWorker(QThread):
    progress = Signal(int)
    finished = Signal(object)
    error = Signal(str)

    def __init__(self, task_type, file_path, data=None):
        super().__init__()
        self.task_type = task_type
        self.file_path = file_path
        self.data = data

    def run(self):
        try:
            if self.task_type == 'export':
                with open(self.file_path, 'w', encoding='utf-8') as f:
                    json.dump(self.data, f, indent=4)
                self.progress.emit(100)
                self.finished.emit(True)

            elif self.task_type == 'import':
                with open(self.file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                self.progress.emit(100)
                self.finished.emit(data)
        except Exception as e:
            self.error.emit(str(e))

# ---------------- Redis Store -----------------
class PolicyStore:
    def __init__(self):
        self.r = None

    def connect(self, host, port):
        self.r = redis.Redis(host=host, port=int(port), decode_responses=True)
        self.r.ping()

    def list_domains(self):
        if not self.r: return []
        dom=set()
        for key in self.r.scan_iter("domain:*"):
            dom.add(key.split(":",1)[1])
        return sorted(dom)

    def load_domain(self, domain):
        raw=self.r.get(f"domain:{domain}")
        if not raw:
            return {"domain":domain,"items":[]}
        data=json.loads(raw)
        for item in data.get("items",[]):
            if "id" not in item:
                item["id"]=str(uuid.uuid4())
            if not item.get("applies_on"):
                item["applies_on"] = ["-"]
        return data

    def save_domain(self,data):
        self.r.set(f"domain:{data['domain']}",json.dumps(data))

    def delete_domain(self,domain):
        self.r.delete(f"domain:{domain}")

    def export_to_json(self,filename, domain_list=None):
        all_domains=[]
        domains = domain_list if domain_list else []
        for d in domains:
            all_domains.append(self.load_domain(d))
        with open(filename,"w") as f:
            json.dump(all_domains,f,indent=4)

    def import_from_json(self,filename):
        with open(filename,"r") as f:
            data=json.load(f)
        for d in data:
            self.save_domain(d)

    # ---------------- Users -----------------
    def get_all_users(self):
        if not self.r: return []
        users=[]
        for key in self.r.scan_iter("user:*"):
            uid = key.split("user:")[1]
            users.append(uid)
        return sorted(users)

    def get_user(self, uid):
        if not uid or uid == "-":
            return {"applies_on":["-"], "uid":"-", "id":"-"}
        tmp = self.r.get(f"user:{uid}")
        if not tmp: return None
        data=json.loads(tmp)
        data["applies_on"]=list(data.get("applies_on", []))
        return data

    def get_domain_policies(self, domain):
        tmp = self.r.get(f"domain:{domain}")
        if not tmp: return []
        policies=json.loads(tmp)
        for item in policies.get("items", []):
            for key in ["id","path","wildcard","action","applies_on",
                        "source_net","policy_type","description","users"]:
                item.setdefault(key,None)
            yield item

# ---------------- Redis Worker -----------------
class RedisWorker(QThread):
    domains_loaded = Signal(list)
    domain_data_loaded = Signal(dict)
    error = Signal(str)

    def __init__(self, store, operation, data=None, parent=None):
        super().__init__(parent)
        self.store=store
        self.operation=operation
        self.data=data

    def run(self):
        try:
            if self.operation=="list_domains":
                self.domains_loaded.emit(self.store.list_domains())
            elif self.operation=="load_domain":
                self.domain_data_loaded.emit(self.store.load_domain(self.data))
        except redis.exceptions.ConnectionError:
            self.error.emit("Redis connection error.")
        except Exception as e:
            self.error.emit(str(e))

# ---------------- Domain Item Dialog -----------------
class DomainItemDialog(QDialog):
    def __init__(self,item=None,parent=None):
        super().__init__(parent)
        self.setWindowTitle("Edit Domain Item")
        self.item=item or {}
        lay=QFormLayout(self)

        self.action=QLineEdit(self.item.get("action","allow"))
        self.source=QLineEdit(",".join(self.item.get("source_net",[])))
        self.applies=QLineEdit(",".join(self.item.get("applies_on",[])))
        self.ptype=QLineEdit(self.item.get("policy_type","custom"))
        self.desc=QTextEdit(self.item.get("description",""))
        self.path=QLineEdit(self.item.get("path","/"))
        self.wild=QCheckBox()
        self.wild.setChecked(self.item.get("wildcard",False))

        lay.addRow("Action",self.action)
        lay.addRow("Source Net",self.source)
        lay.addRow("Applies On",self.applies)
        lay.addRow("Policy Type",self.ptype)
        lay.addRow("Description",self.desc)
        lay.addRow("Path",self.path)
        lay.addRow("Wildcard",self.wild)

        btns=QDialogButtonBox(QDialogButtonBox.Ok|QDialogButtonBox.Cancel)
        btns.accepted.connect(self.accept)
        btns.rejected.connect(self.reject)
        lay.addWidget(btns)

    def get_data(self):
        return {
            "id": self.item.get("id",str(uuid.uuid4())),
            "action": self.action.text().strip(),
            "source_net":[x.strip() for x in self.source.text().split(",") if x.strip()],
            "applies_on":[x.strip() for x in self.applies.text().split(",") if x.strip()] or ["-"],
            "policy_type":self.ptype.text().strip(),
            "description":self.desc.toPlainText().strip(),
            "path":self.path.text().strip(),
            "wildcard":self.wild.isChecked(),
        }

# ---------------- Simulation Logic -----------------
def in_network(src, networks):
    if not networks or src=="-": return True
    try: ip=ipaddress.ip_address(src)
    except ValueError: return False
    for net in networks:
        try:
            if ip in ipaddress.ip_network(net): return True
        except ValueError: continue
    return False

def match_policy(store, user, src, uri):
    if "://" not in uri:
        base_domain=uri.split(":")[0]
        path="/"
    else:
        parsed=urlparse(uri)
        base_domain=parsed.netloc.split(":")[0]
        path=parsed.path if parsed.path else "/"

    user_data=store.get_user(user)
    if not user_data:
        return {"message":f"ERR message=\"no (valid) IDENT {user}\"","user":{"applies_on":["-"]}}

    fqdn=base_domain
    matched_policy=None

    while True:
        for policy in store.get_domain_policies(fqdn):
            match_main = path.startswith(policy["path"]) and fqdn==base_domain
            match_parent = policy["path"]=="/" and fqdn!=base_domain and policy.get("wildcard",False)
            if (match_main or match_parent) and set(policy["applies_on"]) & set(user_data["applies_on"]):
                if not in_network(src, policy["source_net"]): continue
                matched_policy=policy
                matched_policy["domain"]=fqdn
                break
        if matched_policy or "." not in fqdn: break
        fqdn=fqdn.split(".",1)[1]

    if not matched_policy:
        for policy in store.get_domain_policies("*"):
            if set(policy["applies_on"]) & set(user_data["applies_on"]):
                if not in_network(src, policy["source_net"]): continue
                matched_policy=policy
                matched_policy["domain"]="*"
                break

    if matched_policy:
        if matched_policy["action"]=="allow":
            message=f"OK message=\"whitelisted {matched_policy.get('id','')}\" user=\"{user}\""
        else:
            message=f"ERR denied {matched_policy.get('id','')} user=\"{user}\""
    else:
        matched_policy={"action":"allow","policy_type":"fallback"}
        message=f"OK message=\"no (valid) IDENT {user}\""

    return {"message":message,"user":user_data,"policy":matched_policy}

# ---------------- Policy GUI -----------------
class PolicyGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Fantastic v5.7 – Enterprise Policy Manager")
        self.resize(1600,1000)
        self.store=PolicyStore()
        self.current_domain=None
        self.worker=None
        self.updating_items=False
        self.staged_domains={}     # staging changes
        self.deleted_domains=set()
        self.highlighted_domains={} # highlight persistence
        self.init_ui()
        self.setup_auto_refresh()

    # ---------------- UI -----------------
    def eventFilter(self, source, event):
        if event.type() == QEvent.Type.KeyPress and source is self.domain_list:
            if event.key() == Qt.Key_A and (event.modifiers() & Qt.ControlModifier):
                # 1. Stop Python-level signals
                self.updating_items = True
                self.domain_list.blockSignals(True)

                # 2. Use the NATIVE C++ method (Instant)
                self.domain_list.selectAll()

                # 3. Quickly sync checkboxes (the only loop we need)
                for i in range(self.domain_list.count()):
                    self.domain_list.item(i).setCheckState(Qt.CheckState.Checked)

                # 4. Clean up
                self.domain_list.blockSignals(False)
                self.updating_items = False
                self.domain_list.viewport().update()
                return True
        return super().eventFilter(source, event)

    def show_context_menu(self, position):
        menu = QMenu(self)
        sel_all = menu.addAction("Select All")
        desel_all = menu.addAction("Deselect All")

        action = menu.exec(self.domain_list.mapToGlobal(position))
        if not action:
            return

        # 1. Complete freeze of signals and selection logic
        self.domain_list.blockSignals(True)
        self.updating_items = True

        # 2. Temporary change selection mode to stop internal calculation
        old_mode = self.domain_list.selectionMode()
        self.domain_list.setSelectionMode(QListWidget.SelectionMode.NoSelection)

        try:
            if action == sel_all:
                for i in range(self.domain_list.count()):
                    itm = self.domain_list.item(i)
                    itm.setCheckState(Qt.CheckState.Checked)
                    itm.setSelected(True)

            elif action == desel_all:
                for i in range(self.domain_list.count()):
                    itm = self.domain_list.item(i)
                    itm.setCheckState(Qt.CheckState.Unchecked)
                    itm.setSelected(False)
        finally:
            # 3. Restore everything
            self.domain_list.setSelectionMode(old_mode)
            self.updating_items = False
            self.domain_list.blockSignals(False)
            self.domain_list.viewport().update()

        # 2. Re-enable signals and update the UI
        self.updating_items = False
        self.domain_list.blockSignals(False)
        self.domain_list.viewport().update()

    def init_ui(self):
        splitter=QSplitter(Qt.Horizontal)
        self.setCentralWidget(splitter)
        left=QWidget(); l=QVBoxLayout(left)

        # Redis connect
        conn=QHBoxLayout()
        self.host=QLineEdit("127.0.0.1")
        self.port=QLineEdit("6379")
        btn_conn=QPushButton("Connect")
        btn_conn.clicked.connect(self.redis_connect)
        self.ind=QLabel(); self.ind.setFixedSize(16,16)
        self.update_ind(False)
        btn_refresh=QPushButton("Refresh Redis"); btn_refresh.clicked.connect(self.reload_all)
        conn.addWidget(self.host); conn.addWidget(self.port); conn.addWidget(btn_conn)
        conn.addWidget(btn_refresh); conn.addWidget(self.ind)
        l.addLayout(conn)

        # Domain list
        l.addWidget(QLabel("Domains"))
        self.domain_list=QListWidget()
        self.domain_list.setSelectionMode(QListWidget.MultiSelection)
        self.domain_list=QListWidget()
        self.domain_list.setSelectionMode(QListWidget.MultiSelection)


        self.domain_list.setContextMenuPolicy(Qt.CustomContextMenu)
        self.domain_list.customContextMenuRequested.connect(self.show_context_menu)
        self.domain_list.installEventFilter(self)


        self.domain_list.itemClicked.connect(self.load_domain)
        self.domain_list.itemClicked.connect(self.load_domain)
        self.domain_list.itemChanged.connect(self.domain_item_changed)
        self.domain_list.selectionModel().selectionChanged.connect(self.domain_selection_changed)
        l.addWidget(self.domain_list)

        # CRUD Buttons
        db=QHBoxLayout()
        for t,f in[("Add",self.add_domain),("Edit",self.edit_domain),
                   ("Delete",self.del_domain),("Clone",self.clone_domain)]:
            b=QPushButton(t); b.clicked.connect(lambda checked,func=f:self.run_async(func)); db.addWidget(b)
        l.addLayout(db)

        # Domain Items
        l.addWidget(QLabel("Domain Rules"))
        self.item_list=QListWidget(); l.addWidget(self.item_list)
        ib=QHBoxLayout()
        for t,f in[("Add",self.add_item),("Edit",self.edit_item),
                   ("Delete",self.del_item),("Clone",self.clone_item)]:
            b=QPushButton(t); b.clicked.connect(f); ib.addWidget(b)
        l.addLayout(ib)

        # Simulator
        l.addWidget(QLabel("Policy Simulator"))
        sim_fields = QHBoxLayout()
        self.user_sim = QComboBox(); self.user_sim.setFixedWidth(200)
        self.ip_sim = QLineEdit(); self.ip_sim.setPlaceholderText("IP Address"); self.ip_sim.setFixedWidth(150)
        self.domain_sim = QLineEdit(); self.domain_sim.setPlaceholderText("Domain/Path (e.g., cnn.com)"); self.domain_sim.setMinimumWidth(300)
        sim_fields.addWidget(self.user_sim); sim_fields.addWidget(self.ip_sim); sim_fields.addWidget(self.domain_sim)
        sim_fields.setSpacing(10); sim_fields.setContentsMargins(0, 0, 0, 0)
        l.addLayout(sim_fields)

        btn_sim = QPushButton("Run Simulator"); btn_sim.clicked.connect(self.run_simulator)
        l.addWidget(btn_sim)

        l.addWidget(QLabel("Simulator Log"))
        self.log_console=QTextEdit(); self.log_console.setReadOnly(True); l.addWidget(self.log_console)

        self.btn_apply = QPushButton("Apply Changes"); self.btn_apply.setEnabled(False)
        self.btn_apply.clicked.connect(self.apply_changes)
        l.addWidget(self.btn_apply)

        exim=QHBoxLayout()
        btn_exp=QPushButton("Export JSON"); btn_exp.clicked.connect(self.export_json)
        btn_imp=QPushButton("Import JSON"); btn_imp.clicked.connect(self.import_json)
        exim.addWidget(btn_exp); exim.addWidget(btn_imp); l.addLayout(exim)

        self.all_buttons = [b for b in left.findChildren(QPushButton) if b.text() != "Connect"]
        for b in self.all_buttons: b.setEnabled(False)

        splitter.addWidget(left)

    # ---------------- Redis Connect -----------------
    def redis_connect(self):
        try:
            self.store.connect(self.host.text(), self.port.text())
            self.update_ind(True)
            for b in self.all_buttons: b.setEnabled(True)
            QTimer.singleShot(100, self.populate_user_dropdown)
            QTimer.singleShot(100, lambda: self.start_worker("list_domains"))
        except Exception:
            self.update_ind(False)

    def populate_user_dropdown(self):
        self.user_sim.clear(); self.user_sim.addItem("-")
        for uid in self.store.get_all_users(): self.user_sim.addItem(uid)

    def update_ind(self,status):
        pal=self.ind.palette()
        pal.setColor(QPalette.Window,QColor("green" if status else "red"))
        self.ind.setAutoFillBackground(True)
        self.ind.setPalette(pal)

    def start_worker(self,op,data=None):
        self.worker=RedisWorker(self.store,op,data,self)
        self.worker.domains_loaded.connect(self.domains_loaded)
        self.worker.domain_data_loaded.connect(self.domain_loaded)
        self.worker.error.connect(lambda msg: None)
        self.worker.start()

    def run_async(self,func):
        try:
            func()
            self.start_worker("list_domains")
        except Exception: pass

    # ---------------- Domain Load -----------------
    def load_domain(self,item): self.start_worker("load_domain",item.text())
    def domain_loaded(self,data):
        self.current_domain=data
        self.reload_items()

    # ---------------- Domain CRUD with Staging -----------------
    def add_domain(self):
        name, ok = QInputDialog.getText(self, "Add Domain", "Domain:")
        if ok and name:
            name = name.strip()
            if name in self.staged_domains or name in self.store.list_domains(): return
            self.staged_domains[name] = {"domain": name, "items": []}
            self.mark_modified() # Let the worker draw the yellow item

    def edit_domain(self):
        sel = self.domain_list.currentItem()
        if not sel: return
        old = sel.text()
        new, ok = QInputDialog.getText(self, "Edit Domain", "Domain:", text=old)

        if ok and new and old != new:
            new = new.strip()
            # 1. Move data to new name (Yellow)
            data = self.staged_domains.pop(old, self.store.load_domain(old))
            data['domain'] = new
            self.staged_domains[new] = data

            # 2. Mark old name as deleted (Red)
            # Only if 'old' actually exists in Redis
            if old in self.store.list_domains():
                self.deleted_domains.add(old)

            self.mark_modified()

    def del_domain(self):
        selected = [i for i in self.domain_list.selectedItems()]
        if not selected: return

        for item in selected:
            name = item.text()
            # If it's a brand new (unsaved) domain, just remove it entirely
            if name in self.staged_domains and name not in self.store.list_domains():
                self.staged_domains.pop(name, None)
            else:
                # Otherwise, move it to the deleted (Red) set
                self.staged_domains.pop(name, None)
                self.deleted_domains.add(name)

        self.mark_modified() # This triggers the Red paint

    def clone_domain(self):
        sel = self.domain_list.currentItem()
        if not sel: return
        dom = sel.text()
        data = self.staged_domains.get(dom, self.store.load_domain(dom)).copy()
        new_name = dom + "_clone"
        data['domain'] = new_name
        self.staged_domains[new_name] = data
        self.mark_modified() # Let the worker draw the yellow clone

    # ---------------- Domain Items CRUD -----------------
    def add_item(self):
        if not self.current_domain: return
        dlg=DomainItemDialog(parent=self)
        if dlg.exec():
            itm = dlg.get_data()
            self.current_domain['items'].append(itm)
            self.staged_domains[self.current_domain['domain']] = self.current_domain
            self.mark_modified()
            self.reload_items(highlight_id=itm['id'], color="yellow")

    def edit_item(self):
        sel=self.item_list.currentItem()
        if not sel or not self.current_domain: return
        iid=sel.data(Qt.UserRole)
        itm=[i for i in self.current_domain['items'] if i['id']==iid][0]
        dlg=DomainItemDialog(itm,self)
        if dlg.exec():
            new=dlg.get_data()
            idx=self.current_domain['items'].index(itm)
            self.current_domain['items'][idx] = new
            self.staged_domains[self.current_domain['domain']] = self.current_domain
            self.mark_modified()
            self.reload_items(highlight_id=new['id'], color="yellow")

    def del_item(self):
        sel = self.item_list.currentItem()
        if not sel or not self.current_domain: return
        if QMessageBox.question(self,"Confirm Delete","Delete selected rule?",QMessageBox.Yes|QMessageBox.No)!=QMessageBox.Yes: return
        iid = sel.data(Qt.UserRole)
        self.current_domain['items'] = [i for i in self.current_domain['items'] if i['id'] != iid]
        self.staged_domains[self.current_domain['domain']] = self.current_domain
        self.mark_modified()
        self.reload_items(highlight_id=iid, color="red")

    def clone_item(self):
        sel=self.item_list.currentItem()
        if not sel or not self.current_domain: return
        iid=sel.data(Qt.UserRole)
        itm=[i for i in self.current_domain['items'] if i['id']==iid][0]
        new=itm.copy(); new['id']=str(uuid.uuid4())
        self.current_domain['items'].append(new)
        self.staged_domains[self.current_domain['domain']] = self.current_domain
        self.mark_modified()
        self.reload_items(highlight_id=new['id'], color="yellow")

    # ---------------- Reload & Highlight -----------------
    def reload_items(self, highlight_id=None, color=None):
        self.item_list.clear()
        if not self.current_domain: return
        for itm in self.current_domain['items']:
            i=QListWidgetItem(itm['action']+" | "+",".join(itm['source_net']))
            i.setData(Qt.UserRole,itm['id'])
            if highlight_id and itm['id']==highlight_id: i.setBackground(QColor(color))
            self.item_list.addItem(i)

    def add_domain_list_item(self, name, color="yellow"):
        item=QListWidgetItem(name)
        item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
        item.setCheckState(Qt.Unchecked)
        item.setBackground(QColor(color))
        self.highlighted_domains[name] = color
        self.domain_list.addItem(item)

    def highlight_item(self, item, color):
        item.setBackground(QColor(color))
        self.highlighted_domains[item.text()] = color

    def mark_modified(self):
        self.btn_apply.setEnabled(True)
        # Force the worker to fetch and the UI to "re-paint" colors
        self.start_worker("list_domains")

    # ---------------- Apply Changes -----------------
    def apply_changes(self):
        for name, data in self.staged_domains.items():
            self.store.save_domain(data)
        for name in self.deleted_domains:
            self.store.delete_domain(name)

        # Clear staging dictionaries to reset highlights to normal
        self.staged_domains.clear()
        self.deleted_domains.clear()

        self.btn_apply.setEnabled(False)
        self.start_worker("list_domains")

    # ---------------- Simulator & Refresh -----------------
    def run_simulator(self):
        self.log_console.clear()
        user=self.user_sim.currentText()
        ip=self.ip_sim.text().strip() or "-"
        dom_path=self.domain_sim.text().strip()
        if not dom_path: self.log_console.setPlainText("Enter domain/path to simulate"); return
        result=match_policy(self.store,user,ip,dom_path)
        self.log_console.setPlainText(json.dumps(result,indent=2,ensure_ascii=False))

    def setup_auto_refresh(self):
        self.timer = QTimer()
        self.timer.timeout.connect(self.reload_all)
        self.timer.start(5000)

    def domain_item_changed(self, item):
        if self.updating_items: return

        self.updating_items = True
        # Sync the highlight to the checkmark state
        item.setSelected(item.checkState() == Qt.CheckState.Checked)
        self.updating_items = False

    def domain_selection_changed(self, selected, deselected):
        if self.updating_items: return

        self.updating_items = True
        # Block signals so checking boxes doesn't trigger 'domain_item_changed'
        self.domain_list.blockSignals(True)

        try:
            for idx in range(self.domain_list.count()):
                itm = self.domain_list.item(idx)
                # Sync the checkmark to the highlight state
                itm.setCheckState(Qt.CheckState.Checked if itm.isSelected() else Qt.CheckState.Unchecked)
        finally:
            self.domain_list.blockSignals(False)
            self.updating_items = False

    # ---------------- Reload Domains -----------------
    def reload_all(self):
        if self.domain_list.selectedItems() or self.staged_domains or self.deleted_domains or self.highlighted_domains:
            return  # pause refresh when highlighted/staged/deleted items exist
        self.start_worker("list_domains")

    def domains_loaded(self, domains):
        self.domain_list.clear()

        # Combine everything: Redis items + Staged items + Deleted items
        all_view = sorted(list(set(domains) | set(self.staged_domains.keys()) | self.deleted_domains))

        for d in all_view:
            item = QListWidgetItem(d)
            item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
            item.setCheckState(Qt.Unchecked)

            # Priority 1: If it's in deleted_domains, paint it RED
            if d in self.deleted_domains:
                item.setBackground(QColor("red"))
            # Priority 2: If it's in staged_domains, paint it YELLOW
            elif d in self.staged_domains:
                item.setBackground(QColor("yellow"))

            self.domain_list.addItem(item)

    # ---------------- Import/Export -----------------
    def export_json(self):
        file_path, _ = QFileDialog.getSaveFileName(self, "Export Domains", "", "JSON Files (*.json)")
        if not file_path: return

        # --- Your Original Data Preparation Logic ---
        export_data = {}
        all_names = self.store.list_domains()
        for name in all_names:
            if name in self.deleted_domains: continue
            export_data[name] = self.staged_domains.get(name, self.store.load_domain(name))

        for name, data in self.staged_domains.items():
            if name not in export_data: export_data[name] = data

        # --- Progress Bar Setup ---
        self.export_pd = QProgressDialog("Exporting...", "Cancel", 0, 100, self)
        self.export_pd.setMinimumDuration(0)
        self.export_pd.setValue(50) # Set to 50% immediately since preparation is done

        self.ex_worker = JsonWorker('export', file_path, export_data)
        self.ex_worker.finished.connect(lambda: (self.export_pd.close(),
                                                QMessageBox.information(self, "Success", "Exported successfully.")))
        self.ex_worker.error.connect(lambda e: (self.export_pd.close(),
                                               QMessageBox.critical(self, "Error", e)))
        self.ex_worker.start()

    def import_json(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Import Domains", "", "JSON Files (*.json)")
        if not file_path: return

        self.import_pd = QProgressDialog("Reading File...", "Cancel", 0, 100, self)
        self.import_pd.setMinimumDuration(0)
        self.import_pd.setValue(10)

        self.im_worker = JsonWorker('import', file_path)

        def finalize_import(new_data):
            self.import_pd.setValue(80)
            if not isinstance(new_data, dict):
                self.import_pd.close()
                QMessageBox.critical(self, "Error", "Invalid JSON format.")
                return

            # --- Your Original Merge Logic ---
            count = 0
            for domain_name, domain_content in new_data.items():
                self.staged_domains[domain_name] = domain_content
                if domain_name in self.deleted_domains:
                    self.deleted_domains.remove(domain_name)
                count += 1

            self.import_pd.setValue(100)
            self.import_pd.close()
            self.mark_modified() # Triggers Painter
            QMessageBox.information(self, "Success", f"Imported {count} domains.")

        self.im_worker.finished.connect(finalize_import)
        self.im_worker.error.connect(lambda e: (self.import_pd.close(), QMessageBox.critical(self, "Error", e)))
        self.im_worker.start()

# ---------------- Run -----------------
if __name__=="__main__":
    app=QApplication(sys.argv)
    win=PolicyGUI()
    win.show()
    sys.exit(app.exec())