main.py

import os
import sys
import time
import json
import threading
import random
import string
from datetime import datetime, timedelta
import tkinter as tk
from tkinter import messagebox, simpledialog, scrolledtext

# ---------------- Configuration ----------------
if os.name == 'nt':  # Windows
    HOSTS_PATH = r"C:\Windows\System32\drivers\etc\hosts"
else:
    HOSTS_PATH = "/etc/hosts"

REDIRECT_IPS = ["127.0.0.1", "0.0.0.0"]

CONFIG_FILE = "config.json"
OTP_FILE = "otps.txt"

# Block list files for each profile.
BLOCK_LIST_FILES = {
    "study": "block_study.txt",
    "freetime": "block_freetime.txt"
}

# ---------------- Global State ----------------
# Profiles now include per-day schedules: each day 0-6 has active,start,end
state = {
    "profile": "study",
    "profiles": {
        "study": {
            # schedule: day -> {active, start, end}
            "schedule": {i: {"active": (i < 6), "start": "09:00", "end": "17:00"} for i in range(7)}
        },
        "freetime": {
            "schedule": {i: {"active": True, "start": "12:00", "end": "20:00"} for i in range(7)}
        }
    },
    "override_active": False,
    "override_expiry": None,
    "otps": []
}

# ---------------- Configuration Persistence ----------------
def load_config():
    if os.path.isfile(CONFIG_FILE):
        try:
            with open(CONFIG_FILE, "r") as f:
                data = json.load(f)
                state.update({k: data[k] for k in ["profile"] if k in data})
                if "profiles" in data:
                    state["profiles"] = data["profiles"]
        except Exception as e:
            print("Error reading configuration:", e)

def save_config():
    data = {
        "profile": state["profile"],
        "profiles": state["profiles"]
    }
    try:
        with open(CONFIG_FILE, "w") as f:
            json.dump(data, f, indent=4)
    except Exception as e:
        print("Error writing configuration:", e)

load_config()

# ---------------- Block List File Handling ----------------
def load_block_list():
    filename = BLOCK_LIST_FILES[state["profile"]]
    if os.path.isfile(filename):
        try:
            with open(filename, "r") as f:
                state["websites"] = {line.strip() for line in f if line.strip()}
        except Exception as e:
            print(f"Error reading {filename}:", e)
    else:
        state["websites"] = set()

def save_block_list():
    filename = BLOCK_LIST_FILES[state["profile"]]
    try:
        with open(filename, "w") as f:
            for w in state["websites"]:
                f.write(w + "\n")
    except Exception as e:
        print(f"Error writing {filename}:", e)

load_block_list()

# ---------------- OTP Handling ----------------
def load_otps():
    if os.path.isfile(OTP_FILE):
        try:
            with open(OTP_FILE, "r") as f:
                state["otps"] = [line.strip() for line in f if line.strip()]
        except Exception as e:
            print("Error reading OTP file:", e)
    else:
        state["otps"] = []

def save_otp_list():
    try:
        with open(OTP_FILE, "w") as f:
            for otp in state["otps"]:
                f.write(otp + "\n")
    except Exception as e:
        print("Error writing OTP file:", e)

load_otps()

# ---------------- Domain Variations ----------------
def generate_variations(website):
    domain = website.strip().replace("http://", "").replace("https://", "")
    if domain.startswith("www."):
        domain = domain[4:]
    bases = [domain, "www." + domain]
    protos = []
    for b in bases:
        protos += ["http://" + b, "https://" + b]
    return bases + protos

# ---------------- Hosts File Modification ----------------
def update_hosts(block=True):
    try:
        with open(HOSTS_PATH, 'r+') as file:
            lines = file.readlines()
            file.seek(0)
            header, footer = "# Begin Website Blocker\n", "# End Website Blocker\n"
            if block:
                new_lines, in_block = [], False
                for l in lines:
                    if l == header: in_block=True; continue
                    if l == footer: in_block=False; continue
                    if not in_block: new_lines.append(l)
                new_lines.append(header)
                for w in state["websites"]:
                    for v in generate_variations(w):
                        for ip in REDIRECT_IPS:
                            new_lines.append(f"{ip} {v}\n")
                new_lines.append(footer)
                file.truncate(0); file.seek(0); file.writelines(new_lines)
            else:
                new_lines, skipping = [], False
                for l in lines:
                    if l == header: skipping=True; continue
                    if l == footer: skipping=False; continue
                    if not skipping: new_lines.append(l)
                file.truncate(0); file.seek(0); file.writelines(new_lines)
    except Exception as e:
        print("Error updating hosts file:", e)

# ---------------- Blocking Monitor ----------------
def is_currently_blocking():
    now = datetime.now(); d = now.weekday()
    prof = state["profiles"][state["profile"]]["schedule"][str(d)]
    in_range = False
    try:
        start = datetime.strptime(prof["start"], "%H:%M").time()
        end = datetime.strptime(prof["end"], "%H:%M").time()
        t = now.time()
        in_range = start <= t <= end if start <= end else t >= start or t <= end
    except:
        pass
    return prof.get("active", False) and in_range and not state["override_active"]

def monitor_blocker():
    while True:
        now = datetime.now()
        if state["override_active"] and state["override_expiry"] <= now:
            state["override_active"] = False
            state["override_expiry"] = None
        if is_currently_blocking(): update_hosts(True)
        else: update_hosts(False)
        time.sleep(30)

# ---------------- GUI ----------------
class GUI:
    def __init__(self, m):
        self.m = m; m.title("Website Blocker")
        # profile
        pf = tk.LabelFrame(m, text="Profile"); pf.pack(fill="x",padx=5,pady=3)
        self.pv = tk.StringVar(value=state["profile"])
        for i,v in enumerate(["study","freetime"]):
            tk.Radiobutton(pf,text=v.title(),value=v,variable=self.pv,command=self.switch).grid(row=0,column=i)
        # websites
        ws = tk.LabelFrame(m, text="Block List"); ws.pack(fill="both",padx=5,pady=3,expand=True)
        self.txt = scrolledtext.ScrolledText(ws,height=6)
        self.txt.pack(fill="both",expand=True)
        self.txt.bind("<<Modified>>",self.on_ws)
        # schedule table
        sch = tk.LabelFrame(m,text="Weekly Schedule")
        sch.pack(fill="x",padx=5,pady=3)
        days = ["Mon","Tue","Wed","Thu","Fri","Sat","Sun"]
        self.vars = {}
        for idx,d in enumerate(days):
            a = state["profiles"][state["profile"]]["schedule"][str(idx)]["active"]
            v = tk.IntVar(value=1 if a else 0)
            cb= tk.Checkbutton(sch,text=d,variable=v)
            cb.grid(row=0,column=idx)
            self.vars[idx] = {"active":v,
                              "start":tk.Entry(sch,width=5),
                              "end":tk.Entry(sch,width=5)}
            self.vars[idx]["start"].grid(row=1,column=idx); self.vars[idx]["end"].grid(row=2,column=idx)
            s = state["profiles"][state["profile"]]["schedule"][str(idx)]
            self.vars[idx]["start"].insert(0,s["start"])
            self.vars[idx]["end"].insert(0,s["end"])
        tk.Button(m,text="Save Schedule",command=self.save_schedule).pack(pady=3)
        # override
        tk.Button(m,text="Override",command=self.ov).pack(pady=3)
        self.load_ws()
    def switch(self):
        state["profile"]=self.pv.get(); save_config(); load_block_list(); self.load_ws(); self.reload_schedule()
    def load_ws(self):
        self.txt.delete('1.0',tk.END)
        for w in state.get("websites",[]): self.txt.insert(tk.END,w+"\n")
        self.txt.edit_modified(False)
    def on_ws(self,e):
        state["websites"]={l.strip() for l in self.txt.get('1.0',tk.END).splitlines() if l.strip()}
        save_block_list(); self.txt.edit_modified(False)
    def save_schedule(self):
        if is_currently_blocking() and not state["override_active"]:
            messagebox.showerror("Error","Cannot update during active block."); return
        sched=state["profiles"][state["profile"]]["schedule"]
        for i,v in self.vars.items():
            sched[str(i)]["active"]=bool(v["active"].get())
            sched[str(i)]["start"]=v["start"].get()
            sched[str(i)]["end"]=v["end"].get()
        save_config(); print("Schedule saved.")
    def reload_schedule(self):
        sched=state["profiles"][state["profile"]]["schedule"]
        for i,v in self.vars.items():
            v["active"].set(1 if sched[str(i)]["active"] else 0)
            v["start"].delete(0,tk.END); v["start"].insert(0,sched[str(i)]["start"])
            v["end"].delete(0,tk.END); v["end"].insert(0,sched[str(i)]["end"])
    def ov(self):
        otp=simpledialog.askstring("OTP","Enter one-time password:")
        if otp in state["otps"]:
            state["otps"].remove(otp); save_otp_list()
            now=datetime.now(); t=now+timedelta(days=1)
            state["override_active"]=True; state["override_expiry"]=datetime(t.year,t.month,t.day)
            print("Override active until midnight.")
        else: messagebox.showerror("Invalid","Bad OTP.")

def main():
    th=threading.Thread(target=monitor_blocker,daemon=True); th.start()
    root=tk.Tk(); GUI(root); root.mainloop()

if __name__=="__main__": main()