Update bot.py

bot.py

@@ -6,10 +6,14 @@ import asyncio
 from pyrogram import Client, filters
 from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton
 from pyrogram.errors import MessageNotModified
-import db
+
+from state import state
+from logger import get_logger
 import scraper
 from curl_cffi.requests import AsyncSession
 
+log = get_logger()
+
 # --- CONFIG & AUTH ---
 API_ID = int(os.environ.get("API_ID", 0))
 API_HASH = os.environ.get("API_HASH", "")
@@ -17,47 +21,47 @@ BOT_TOKEN = os.environ.get("BOT_TOKEN", "")
 ADMIN_IDS = [int(x) for x in os.environ.get("ADMIN_IDS", "").split(",") if x]
 
 app = Client("bot_session", api_id=API_ID, api_hash=API_HASH, bot_token=BOT_TOKEN)
-
-# Custom admin filter for Pyrogram
 is_admin = filters.user(ADMIN_IDS)
 
-# --- NAVIGATION MENUS ---
-def get_main_menu():
-    return InlineKeyboardMarkup([
-        [InlineKeyboardButton("▶️ Start", callback_data="start_scan"), 
-         InlineKeyboardButton("⏸ Pause", callback_data="pause_scan"), 
-         InlineKeyboardButton("⏹ Stop", callback_data="stop_scan")],
-        [InlineKeyboardButton("📊 Live Status", callback_data="show_status"), 
-         InlineKeyboardButton("📥 Export Data", callback_data="export_files")],
-        [InlineKeyboardButton("⚙️ System Settings", callback_data="menu_settings")]
-    ])
+# --- DYNAMIC UI MENUS ---
+def get_dynamic_menu():
+    """Generates toggle switches based on the current in-memory state."""
+    buttons = []
+    
+    if state.status == "STOPPED":
+        buttons.append([InlineKeyboardButton("▶️ Start Scan", callback_data="toggle_start")])
+    elif state.status == "RUNNING":
+        buttons.append([
+            InlineKeyboardButton("⏸ Pause", callback_data="toggle_pause"),
+            InlineKeyboardButton("⏹ Stop", callback_data="toggle_stop")
+        ])
+    elif state.status == "PAUSED":
+        buttons.append([
+            InlineKeyboardButton("▶️ Resume", callback_data="toggle_resume"),
+            InlineKeyboardButton("⏹ Stop", callback_data="toggle_stop")
+        ])
+
+    buttons.append([InlineKeyboardButton("📊 Live Status", callback_data="show_status"), InlineKeyboardButton("📥 Export Data", callback_data="export_files")])
+    buttons.append([InlineKeyboardButton("⚙️ System Settings", callback_data="menu_settings")])
+    
+    return InlineKeyboardMarkup(buttons)
 
 def get_settings_menu():
     return InlineKeyboardMarkup([
-        [InlineKeyboardButton("💾 Load Words", callback_data="load_words"), 
-         InlineKeyboardButton("📏 Set Length", callback_data="set_min_length")],
-        [InlineKeyboardButton("⚡ Set Speed", callback_data="set_speed"), 
-         InlineKeyboardButton("🔴 Wipe DB", callback_data="reset_confirm")],
+        [InlineKeyboardButton("💾 Load Words", callback_data="load_words"), InlineKeyboardButton("📏 Set Length", callback_data="set_min_length")],
+        [InlineKeyboardButton("⚡ Set Speed", callback_data="set_speed"), InlineKeyboardButton("🔴 Clear Memory", callback_data="reset_confirm")],
         [InlineKeyboardButton("« Back to Main", callback_data="back_main")]
     ])
 
-async def generate_status_msg():
-    state = await db.get_state()
-    counts = await db.get_counts()
-    concurrency = await db.get_concurrency()
-    qlen = await db.queue_size()
-    proxy_count = await db.get_proxy_count()
-    
-    total = int(state.get("total", 0))
-    processed = int(state.get("processed", 0))
-    min_len = int(state.get("min_length", 4))
-    pct = (processed / total * 100) if total > 0 else 0
+def generate_status_msg():
+    qlen = state.queue.qsize() if state.queue else 0
+    pct = (state.processed / state.total_words * 100) if state.total_words > 0 else 0
     
-    run_str = "🟢 RUNNING" if state.get("running") == "1" else "⏹ STOPPED"
-    if state.get("paused") == "1": run_str = "⏸ PAUSED"
+    run_str = f"🟢 {state.status}" if state.status == "RUNNING" else f"🔴 {state.status}"
+    if state.status == "PAUSED": run_str = "⏸ PAUSED"
     
     bar_len = 20
-    filled = int(bar_len * (processed / total)) if total > 0 else 0
+    filled = int(bar_len * (state.processed / state.total_words)) if state.total_words > 0 else 0
     bar = "=" * filled + "-" * (bar_len - filled)
     
     return (
@@ -65,20 +69,19 @@ async def generate_status_msg():
         f"━━━━━━━━━━━━━━━━━━━━━\n"
         f"**State:** `{run_str}`\n"
         f"**Queue:** `{qlen:,}` waiting\n"
-        f"**Speed:** `{concurrency}` workers\n"
-        f"**Filter:** `Min {min_len} letters`\n"
-        f"**Proxies:** `{proxy_count}` Active\n"
+        f"**Speed:** `{state.concurrency}` workers\n"
+        f"**Proxies:** `{len(state.proxies)}` Active\n"
         f"**Progress:** `{pct:.2f}%`\n"
         f"`[{bar}]`\n"
-        f"`{processed:,} / {total:,}`\n\n"
-        f"**DATABASE METRICS**\n"
+        f"`{state.processed:,} / {state.total_words:,}`\n\n"
+        f"**METRICS**\n"
         f"━━━━━━━━━━━━━━━━━━━━━\n"
-        f"���� Taken      : `{counts.get('taken', 0):,}`\n"
-        f"🚫 Unavail    : `{counts.get('unavailable', 0):,}`\n"
-        f"💰 For Sale   : `{counts.get('forsale', 0):,}`\n"
-        f"🔨 Auction    : `{counts.get('auction', 0):,}`\n"
-        f"✅ Available  : `{counts.get('available', 0):,}`\n"
-        f"🛒 Sold       : `{counts.get('sold', 0):,}`\n"
+        f"🔴 Taken      : `{state.counts.get('taken', 0):,}`\n"
+        f"🚫 Unavail    : `{state.counts.get('unavailable', 0):,}`\n"
+        f"💰 For Sale   : `{state.counts.get('forsale', 0):,}`\n"
+        f"🔨 Auction    : `{state.counts.get('auction', 0):,}`\n"
+        f"✅ Available  : `{state.counts.get('available', 0):,}`\n"
+        f"🛒 Sold       : `{state.counts.get('sold', 0):,}`\n"
     )
 
 # --- PROXY TESTER ---
@@ -89,16 +92,15 @@ async def test_proxy(proxy_url):
         async with AsyncSession(impersonate="chrome120", proxies=proxies) as session:
             resp = await session.get("https://fragment.com/", timeout=10)
             if resp.status_code == 200:
-                latency = round((time.time() - start) * 1000)
-                return proxy_url, latency
+                return proxy_url, round((time.time() - start) * 1000)
     except Exception:
         pass
     return proxy_url, None
 
-# --- MESSAGE HANDLERS ---
+# --- COMMANDS ---
 @app.on_message(filters.command("start") & is_admin)
 async def start_cmd(client, message):
-    await message.reply_text("**Fragment Control System**", reply_markup=get_main_menu())
+    await message.reply_text("**Fragment Control System**", reply_markup=get_dynamic_menu())
 
 @app.on_message(filters.command("proxies") & is_admin)
 async def proxy_cmd(client, message):
@@ -119,294 +121,107 @@ async def proxy_cmd(client, message):
     tasks = [test_proxy(p) for p in raw_proxies]
     results = await asyncio.gather(*tasks)
     
-    working_proxies = []
-    latency_board = []
-    
-    for p_url, latency in results:
-        if latency is not None:
-            working_proxies.append(p_url)
-            latency_board.append((p_url, latency))
-            
-    latency_board.sort(key=lambda x: x[1])
-    await db.save_working_proxies(working_proxies)
-    
-    top_str = "\n".join([f"⏱ `{lat}ms` | {url.split('@')[-1] if '@' in url else url}" for url, lat in latency_board[:10]])
-    report = (
-        f"✅ **Proxy Test Complete**\n"
-        f"━━━━━━━━━━━━━━━━\n"
-        f"**Total Tested:** `{len(raw_proxies)}`\n"
-        f"**Working & Saved:** `{len(working_proxies)}`\n\n"
-        f"🏆 **Top Fastest Proxies:**\n{top_str if top_str else 'None'}"
-    )
-    await msg.edit_text(report, reply_markup=get_main_menu())
-
-@app.on_message(filters.command("check") & is_admin)
-async def check_cmd(client, message):
-    if len(message.command) < 2: return await message.reply_text("Usage: `/check <word>`")
-    target = message.command[1].lower()
-    msg = await message.reply_text(f"📸 Taking live screenshot of `@{target}` on Fragment...")
-    
-    url = f"https://fragment.com/username/{target}"
-    img_filename = None
+    working = [p for p, lat in results if lat is not None]
+    state.proxies = working # Save to memory
     
-    try:
-        def take_screenshot():
-            from html2image import Html2Image
-            hti = Html2Image(
-                browser_executable='/usr/bin/chromium',
-                custom_flags=['--no-sandbox', '--disable-gpu', '--hide-scrollbars', '--disable-dev-shm-usage']
-            ) 
-            hti.size = (1000, 750) 
-            name = f"{target}_fragment.png"
-            hti.screenshot(url=url, save_as=name)
-            return name
-
-        img_filename = await asyncio.to_thread(take_screenshot)
-        await message.reply_document(document=img_filename, caption=f"**Live Web Snapshot:** `@{target}`")
-        await msg.delete()
-            
-    except Exception as e:
-        try: await msg.edit_text(f"❌ **Screenshot Error:**\n`{str(e)}`")
-        except: pass
-    finally:
-        if img_filename and os.path.exists(img_filename):
-            os.remove(img_filename)
-
-@app.on_message(filters.command("html") & is_admin)
-async def html_cmd(client, message):
-    if len(message.command) < 2: return await message.reply_text("Usage: `/html <word>`")
-    target = message.command[1].lower()
-    msg = await message.reply_text(f"🔍 Fetching raw HTML for `@{target}` via Chrome Impersonation...")
-    
-    url = f"https://fragment.com/username/{target}"
-    
-    try:
-        async with AsyncSession(impersonate="chrome120") as session:
-            resp = await session.get(url, timeout=15, allow_redirects=True)
-            html_text = resp.text
-            status_code = resp.status_code
-            final_url = str(resp.url)
-                
-        f = io.BytesIO(html_text.encode('utf-8'))
-        f.name = f"{target}_fragment.html"
-        f.seek(0) 
-        
-        caption = (f"📄 **Raw Response for** `@{target}`\n"
-                   f"**Requested URL:** `{url}`\n"
-                   f"**Final URL:** `{final_url}`\n"
-                   f"**Status Code:** `{status_code}`")
-        
-        await message.reply_document(document=f, caption=caption)
-        await msg.delete() 
-    except Exception as e:
-        try: await msg.edit_text(f"❌ **Error fetching HTML:** `{str(e)}`")
-        except: pass
-
-@app.on_message(filters.command("upload") & is_admin)
-async def upload_cmd(client, message):
-    if not message.reply_to_message or not message.reply_to_message.document:
-        return await message.reply_text("Reply to a `.txt` file.")
-    
-    msg = await message.reply_text("Downloading...")
-    await message.reply_to_message.download(file_name="words.txt")
-    await msg.edit_text("✅ Uploaded. Go to Settings -> Load Database.", reply_markup=get_main_menu())
-
-
-# --- STRING SESSION STATE MACHINE ---
-# Pyrogram doesn't have native "conversation" handlers, so we use Futures
-input_waiters = {}
-
-@app.on_message(filters.text & is_admin, group=-1)
-async def catch_input(client, message):
-    chat_id = message.chat.id
-    if chat_id in input_waiters and not input_waiters[chat_id].done():
-        input_waiters[chat_id].set_result(message.text)
-        message.stop_propagation()
-
-async def get_response(chat_id, timeout=300):
-    loop = asyncio.get_running_loop()
-    future = loop.create_future()
-    input_waiters[chat_id] = future
-    try: return await asyncio.wait_for(future, timeout)
-    finally: del input_waiters[chat_id]
+    await msg.edit_text(f"✅ Saved **{len(working)}** working proxies to memory.", reply_markup=get_dynamic_menu())
 
-@app.on_message(filters.command("getstring") & is_admin)
-async def pyrogram_string_cmd(client, message):
-    chat_id = message.chat.id
-    try:
-        await message.reply_text("🛠 **Pyrogram Session Generator**\nEnter your `API_ID`:")
-        api_id = int((await get_response(chat_id)).strip())
-
-        await message.reply_text("Enter your `API_HASH`:")
-        api_hash = (await get_response(chat_id)).strip()
-
-        await message.reply_text("Enter the **Phone Number** (e.g., +1234567890):")
-        phone = (await get_response(chat_id)).replace(" ", "").strip()
-
-        await message.reply_text("⏳ Requesting OTP from Telegram...")
-        
-        from pyrogram import Client as TmpClient
-        from pyrogram.errors import SessionPasswordNeeded
-
-        tmp_app = TmpClient("temp_session", api_id=api_id, api_hash=api_hash, in_memory=True)
-        await tmp_app.connect()
-        
-        sent_code = await tmp_app.send_code(phone)
-        
-        await message.reply_text("📥 **Code Sent!**\n⚠️ Enter code with spaces (e.g., `1 2 3 4 5`):")
-        otp_code = (await get_response(chat_id)).replace(" ", "").strip()
-        
-        try:
-            await tmp_app.sign_in(phone, sent_code.phone_code_hash, otp_code)
-        except SessionPasswordNeeded:
-            await message.reply_text("🔐 **2FA Password Required:**")
-            pwd = (await get_response(chat_id)).strip()
-            await tmp_app.check_password(pwd)
-        
-        string_session = await tmp_app.export_session_string()
-        await message.reply_text(f"✅ **Session Generated!**\n\n`{string_session}`\n\n⚠️ Keep this secret.")
-    except asyncio.TimeoutError:
-        await message.reply_text("❌ **Timeout.** Run `/getstring` again.")
-    except Exception as e:
-        await message.reply_text(f"❌ **Error:** `{str(e)}`")
-    finally:
-        if 'tmp_app' in locals():
-            try: await tmp_app.disconnect()
-            except: pass
-
-# --- CALLBACK ROUTER ---
+# --- UI CALLBACKS ---
 @app.on_callback_query(is_admin)
 async def button_handler(client, query):
     data = query.data
-    
     try:
         if data == "back_main": 
-            await query.edit_message_text("**Fragment Control System**", reply_markup=get_main_menu())
+            await query.edit_message_text("**Fragment Control System**", reply_markup=get_dynamic_menu())
             
         elif data == "menu_settings": 
             await query.edit_message_text("**System Settings**", reply_markup=get_settings_menu())
             
-        elif data == "set_min_length":
-            kb = InlineKeyboardMarkup([
-                [InlineKeyboardButton("4 Letters", callback_data="min_len_4"), InlineKeyboardButton("5 Letters", callback_data="min_len_5")],
-                [InlineKeyboardButton("6 Letters", callback_data="min_len_6"), InlineKeyboardButton("7 Letters", callback_data="min_len_7")],
-                [InlineKeyboardButton("« Back", callback_data="menu_settings")]
-            ])
-            await query.edit_message_text("📏 **Select Minimum Word Length:**", reply_markup=kb)
-            
-        elif data.startswith("min_len_"):
-            new_len = int(data.split("_")[2])
-            await db.set_state(min_length=new_len)
-            await query.edit_message_text(f"✅ Filter updated to `{new_len}` letters.", reply_markup=get_settings_menu())
+        # DYNAMIC TOGGLES
+        elif data == "toggle_start":
+            if not state.queue or state.queue.empty():
+                return await query.answer("Queue is empty! Load words first.", show_alert=True)
+            state.status = "RUNNING"
+            log.info("Scanner Started")
+            await query.edit_message_text("▶️ Scanner Running", reply_markup=get_dynamic_menu())
+            
+        elif data == "toggle_pause":
+            state.status = "PAUSED"
+            log.info("Scanner Paused")
+            await query.edit_message_text("⏸ Scanner Paused", reply_markup=get_dynamic_menu())
+            
+        elif data == "toggle_resume":
+            state.status = "RUNNING"
+            log.info("Scanner Resumed")
+            await query.edit_message_text("▶️ Scanner Resumed", reply_markup=get_dynamic_menu())
+            
+        elif data == "toggle_stop":
+            state.status = "STOPPED"
+            log.info("Scanner Stopped")
+            await query.edit_message_text("⏹ Scanner Stopped", reply_markup=get_dynamic_menu())
             
         elif data == "load_words": 
-            await query.edit_message_text("⏳ Processing and loading chunk...")
+            await query.edit_message_text("⏳ Processing and loading chunk into memory...")
             try:
                 with open("words.txt", "r", encoding="utf-8") as f: 
                     all_lines = [line.strip().lower() for line in f if line.strip()]
             except Exception:
                 return await query.edit_message_text("❌ `words.txt` not found. Upload it first.")
                 
-            if not all_lines: return await query.edit_message_text("Empty file.")
-                
-            state = await db.get_state()
-            min_length = int(state.get("min_length", 4))
-                
             CHUNK = 200000
             current = all_lines[:CHUNK]
-            remaining = all_lines[CHUNK:]
             
-            pattern = re.compile(rf'^[a-z0-9_]{{{min_length},32}}$')
+            pattern = re.compile(rf'^[a-z0-9_]{{{state.min_length},32}}$')
             valid = {w for w in current if pattern.match(w)}
             
-            r = await db.get_redis()
-            done = await r.smembers("frag:done")
-            to_q = list(valid - done)
-            
-            if to_q:
-                pipe = r.pipeline()
-                for i in range(0, len(to_q), 1000): pipe.lpush("frag:queue", *to_q[i:i+1000])
-                await pipe.execute()
-                
-            with open("words.txt", "w", encoding="utf-8") as f: 
-                f.write("\n".join(remaining))
-                
-            await db.set_state(total=len(valid))
-            await query.edit_message_text(f"✅ Loaded {len(to_q):,} words.", reply_markup=get_main_menu())
-            
-        elif data == "start_scan": 
-            await db.set_state(running="1", paused="0")
-            await query.edit_message_text("▶️ Scanner Running", reply_markup=get_main_menu())
-            
-        elif data == "pause_scan":
-            s = await db.get_state()
-            p = "0" if s.get("paused") == "1" else "1"
-            await db.set_state(paused=p)
-            await query.edit_message_text("⏸ Scanner Paused" if p=="1" else "▶️ Resumed", reply_markup=get_main_menu())
-            
-        elif data == "stop_scan":
-            await db.set_state(running="0", paused="0")
-            await query.edit_message_text("⏹ Scanner Stopped", reply_markup=get_main_menu())
+            await state.add_to_queue(list(valid))
+            await query.edit_message_text(f"✅ Loaded {len(valid):,} words into queue.", reply_markup=get_dynamic_menu())
             
         elif data in ("show_status", "refresh_status"):
-            msg = await generate_status_msg()
+            msg = generate_status_msg()
             kb = InlineKeyboardMarkup([[InlineKeyboardButton("🔄 Refresh", callback_data="refresh_status")], [InlineKeyboardButton("« Back", callback_data="back_main")]])
             try: await query.edit_message_text(msg, reply_markup=kb)
             except MessageNotModified: pass
-            
-        elif data == "set_speed":
-            kb = InlineKeyboardMarkup([
-                [InlineKeyboardButton("10", callback_data="spd_10"), InlineKeyboardButton("20", callback_data="spd_20")], 
-                [InlineKeyboardButton("30", callback_data="spd_30"), InlineKeyboardButton("50", callback_data="spd_50")], 
-                [InlineKeyboardButton("« Back", callback_data="menu_settings")]
-            ])
-            await query.edit_message_text("⚡ **Select Concurrent Workers:**", reply_markup=kb)
-            
-        elif data.startswith("spd_"):
-            await db.set_concurrency(int(data.split("_")[1]))
-            await query.edit_message_text("✅ Speed Updated", reply_markup=get_settings_menu())
-            
+
         elif data == "export_files":
-            await query.edit_message_text("⏳ Generating text files...")
-            exported_count = 0
-            
+            await query.edit_message_text("⏳ Preparing files from disk...")
+            exported = 0
             for s in ["taken", "unavailable", "sold", "forsale", "auction", "available"]:
-                lst = sorted(list(await (getattr(db, f"get_all_{s}")())))
-                if lst:
-                    f = io.BytesIO("\n".join(lst).encode('utf-8'))
-                    f.name = f"{s}.txt"
-                    await client.send_document(query.message.chat.id, document=f)
-                    exported_count += 1
-            
-            if os.path.exists("other.txt"):
-                await client.send_document(query.message.chat.id, document="other.txt", caption="Unknown HTML Responses")
-                exported_count += 1
+                if os.path.exists(f"results/{s}.txt"):
+                    await client.send_document(query.message.chat.id, document=f"results/{s}.txt")
+                    exported += 1
                     
-            if exported_count > 0:
-                await query.edit_message_text(f"✅ Export complete! Sent {exported_count} files.", reply_markup=get_main_menu())
+            if exported > 0:
+                await query.edit_message_text(f"✅ Exported {exported} result files.", reply_markup=get_dynamic_menu())
             else:
-                await query.edit_message_text("⚠️ No data to export yet.", reply_markup=get_main_menu())
+                await query.edit_message_text("⚠️ No data files found.", reply_markup=get_dynamic_menu())
                 
-        elif data == "reset_confirm": 
-            kb = InlineKeyboardMarkup([[InlineKeyboardButton("CONFIRM WIPE", callback_data="reset_do")], [InlineKeyboardButton("Cancel", callback_data="menu_settings")]])
-            await query.edit_message_text("⚠️ WIPE ALL REDIS DATA?", reply_markup=kb)
-            
-        elif data == "reset_do":
-            await db.flush_all()
-            await query.edit_message_text("💣 Memory purged successfully.", reply_markup=get_settings_menu())
+        elif data == "reset_confirm":
+            # Reset memory counters
+            state.processed = 0
+            state.total_words = 0
+            state.counts = {k: 0 for k in state.counts}
+            state.queue = asyncio.Queue()
+            state.status = "STOPPED"
+            # Delete output files
+            for file in os.listdir("results"):
+                os.remove(f"results/{file}")
+            await query.edit_message_text("💣 Memory and Results purged.", reply_markup=get_settings_menu())
             
     except Exception as e:
         if not isinstance(e, MessageNotModified):
-            print(f"Callback Error: {e}")
+            log.exception("UI Callback Error:")
 
-# If run directly as a script
 if __name__ == "__main__":
-    from worker import run_worker
+    from worker import manage_workers
     
     async def main_loop():
         await app.start()
-        print("🤖 Pyrogram Bot Started!")
-        asyncio.create_task(run_worker())
+        log.info("🤖 Pyrogram Bot Started!")
+        
+        # Start the worker manager daemon
+        asyncio.create_task(manage_workers())
+        
         from pyrogram import idle
         await idle()
         await app.stop()