dc-securex 2.29.7__py3-none-any.whl

securex/workers/guild_worker.py

@@ -0,0 +1,186 @@
+"""
+Guild Restoration Worker - Restores server name and vanity URL only
+Triggered by GuildHandler when unauthorized changes are detected
+Supports per-guild user tokens with file storage and memory caching
+"""
+import asyncio
+import discord
+import aiohttp
+import json
+from pathlib import Path
+from datetime import datetime
+from typing import Dict, Any, Optional
+
+
+class GuildWorker:
+    """Background worker for guild settings restoration"""
+    
+    def __init__(self, sdk):
+        self.sdk = sdk
+        self.bot = sdk.bot
+        self._worker_task = None
+        
+        # Per-guild user tokens
+        self._user_tokens: Dict[int, str] = {}  # guild_id -> user_token (cache)
+        self._tokens_file = sdk.backup_dir / "user_tokens.json"
+    
+    async def _load_tokens(self):
+        """Load user tokens from file into cache"""
+        try:
+            if self._tokens_file.exists():
+                with open(self._tokens_file, 'r') as f:
+                    data = json.load(f)
+                    # Convert string keys back to int
+                    self._user_tokens = {int(k): v for k, v in data.items()}
+                    print(f"✅ Loaded user tokens for {len(self._user_tokens)} guild(s)")
+        except Exception as e:
+            print(f"Error loading user tokens: {e}")
+    
+    async def _save_tokens(self):
+        """Save user tokens to file"""
+        try:
+            # Convert int keys to string for JSON
+            data = {str(k): v for k, v in self._user_tokens.items()}
+            with open(self._tokens_file, 'w') as f:
+                json.dump(data, f, indent=2)
+        except Exception as e:
+            print(f"Error saving user tokens: {e}")
+    
+    async def set_user_token(self, guild_id: int, token: str):
+        """Set user token for a specific guild"""
+        self._user_tokens[guild_id] = token
+        await self._save_tokens()
+        print(f"✅ User token set for guild {guild_id}")
+    
+    def get_user_token(self, guild_id: int) -> Optional[str]:
+        """Get user token for a guild (from cache)"""
+        return self._user_tokens.get(guild_id)
+    
+    async def remove_user_token(self, guild_id: int):
+        """Remove user token for a guild"""
+        if guild_id in self._user_tokens:
+            del self._user_tokens[guild_id]
+            await self._save_tokens()
+            print(f"❌ User token removed for guild {guild_id}")
+    
+    async def start(self):
+        """Start the guild restoration worker"""
+        # Load tokens first
+        await self._load_tokens()
+        
+        if self._worker_task is None or self._worker_task.done():
+            self._worker_task = asyncio.create_task(self._restoration_loop())
+            print("⚡ Guild Restoration Worker started")
+    
+    async def stop(self):
+        """Stop the worker"""
+        if self._worker_task:
+            self._worker_task.cancel()
+            self._worker_task = None
+    
+    async def _restoration_loop(self):
+        """Main worker loop - processes guild_update entries from queue"""
+        while True:
+            try:
+                entry = await self.sdk.guild_queue.get()
+                
+                # Only process guild_update actions
+                if entry.action != discord.AuditLogAction.guild_update:
+                    continue
+                
+                # Extract changes from entry
+                guild = entry.guild
+                executor = entry.user
+                
+                # Skip if bot did it
+                if executor.id == self.bot.user.id:
+                    continue
+                
+                # Check authorization
+                if executor.id == guild.owner_id:
+                    continue
+                
+                whitelist_set = self.sdk.whitelist_cache.get(guild.id, set())
+                if executor.id in whitelist_set:
+                    continue
+                
+                # Unauthorized change - restore from backup
+                print(f"🔍 Unauthorized guild update by {executor.name}")
+                await self._restore_guild_settings(guild)
+                    
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                print(f"Error in guild restoration worker: {e}")
+    
+    async def _restore_vanity_via_api(self, guild_id: int, vanity_code: str) -> bool:
+        """Restore vanity URL using user token and Discord API"""
+        user_token = self.get_user_token(guild_id)
+        
+        if not user_token:
+            print(f"⚠️ No user token set for guild {guild_id}! Cannot restore vanity URL.")
+            print(f"   Use: await sx.guild_worker.set_user_token({guild_id}, 'USER_TOKEN')")
+            return False
+        
+        try:
+            url = f"https://discord.com/api/v9/guilds/{guild_id}/vanity-url"
+            headers = {
+                "authorization": user_token,
+                "content-type": "application/json"
+            }
+            payload = {"code": vanity_code}
+            
+            async with aiohttp.ClientSession() as session:
+                async with session.patch(url, headers=headers, json=payload) as response:
+                    if response.status == 200:
+                        return True
+                    else:
+                        error_text = await response.text()
+                        print(f"⚠️ Failed to restore vanity: {response.status} - {error_text}")
+                        return False
+        except Exception as e:
+            print(f"Error restoring vanity via API: {e}")
+            return False
+    
+    async def _restore_guild_settings(self, guild):
+        """Restore guild settings (name and vanity) from backup"""
+        try:
+            print(f"🔧 [DEBUG] Starting restoration for guild: {guild.name}")
+            start_time = datetime.now()
+            
+            backup_data = await self.sdk.backup_manager.get_guild_settings(guild.id)
+            print(f"📦 [DEBUG] Backup data retrieved: {backup_data}")
+            if not backup_data:
+                print(f"⚠️ No backup found for guild {guild.name}")
+                return
+            
+            restore_params = {}
+            restored_items = []
+            
+            # Restore vanity URL if we have it in backup
+            if backup_data.get("vanity_url_code"):
+                old_vanity = backup_data["vanity_url_code"]
+                vanity_success = await self._restore_vanity_via_api(guild.id, old_vanity)
+                if vanity_success:
+                    restored_items.append(f"vanity (discord.gg/{old_vanity})")
+            
+            # Restore server name if different from backup
+            if backup_data.get("name") and backup_data["name"] != guild.name:
+                restore_params["name"] = backup_data["name"]
+                restored_items.append(f"name ({backup_data['name']})")
+            
+            if restore_params:
+                await guild.edit(**restore_params, reason="SecureX: Restoring unauthorized guild changes")
+            
+            if restored_items:
+                elapsed = (datetime.now() - start_time).total_seconds() * 1000
+                items_str = ", ".join(restored_items)
+                print(f"⚡ Restored guild settings in {elapsed:.0f}ms: {items_str}")
+            
+        except discord.errors.HTTPException as e:
+            if e.code == 50035:
+                print(f"⚠️ Some settings could not be restored: {e}")
+            else:
+                print(f"Error restoring guild settings: {e}")
+        except Exception as e:
+            print(f"Error in guild restoration: {e}")

securex/workers/log_worker.py

@@ -0,0 +1,88 @@
+"""
+Log Worker - Logs threats and fires callbacks
+"""
+import discord
+import asyncio
+from ..models import ThreatEvent
+
+
+class LogWorker:
+    """Worker that logs threats and triggers callbacks"""
+    
+    def __init__(self, sdk):
+        self.sdk = sdk
+        self.bot = sdk.bot
+        self.log_queue = asyncio.Queue()
+        self._worker_task = None
+        
+        self.action_map = {
+            discord.AuditLogAction.bot_add: "bot_add",
+            discord.AuditLogAction.channel_create: "channel_create",
+            discord.AuditLogAction.channel_delete: "channel_delete",
+            discord.AuditLogAction.channel_update: "channel_update",
+            discord.AuditLogAction.role_create: "role_create",
+            discord.AuditLogAction.role_delete: "role_delete",
+            discord.AuditLogAction.role_update: "role_update",
+            discord.AuditLogAction.kick: "member_kick",
+            discord.AuditLogAction.ban: "member_ban",
+            discord.AuditLogAction.unban: "member_unban",
+            discord.AuditLogAction.member_update: "member_update",
+            discord.AuditLogAction.webhook_create: "webhook_create",
+            discord.AuditLogAction.guild_update: "guild_update",
+        }
+    
+    async def start(self):
+        """Start the log worker"""
+        if self._worker_task is None or self._worker_task.done():
+            self._worker_task = asyncio.create_task(self._worker_loop())
+            print("📝 Log Worker started")
+    
+    async def stop(self):
+        """Stop the worker"""
+        if self._worker_task:
+            self._worker_task.cancel()
+            self._worker_task = None
+    
+    async def _worker_loop(self):
+        """Main worker loop"""
+        while True:
+            try:
+                entry = await self.log_queue.get()
+                await self._process_entry(entry)
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                print(f"Error in log worker: {e}")
+    
+    async def _process_entry(self, entry):
+        """Process log entry and fire callbacks"""
+        try:
+            guild = entry.guild
+            executor = entry.user
+            
+            if executor.id == self.bot.user.id or executor.id == guild.owner_id:
+                return
+            
+            violation_type = self.action_map.get(entry.action)
+            if not violation_type:
+                return
+            
+            guild_punishments = self.sdk.punishment_cache.get(guild.id, self.sdk.punishments)
+            punishment_action = guild_punishments.get(violation_type, "none")
+            
+            threat_event = ThreatEvent(
+                guild_id=guild.id,
+                actor_id=executor.id,
+                type=violation_type,
+                target_id=getattr(entry.target, 'id', None),
+                target_name=getattr(entry.target, 'name', 'Unknown'),
+                prevented=True,
+                restored=True,
+                punishment_action=punishment_action,
+                timestamp=entry.created_at
+            )
+            
+            await self.sdk._trigger_callbacks('threat_detected', threat_event)
+            
+        except Exception as e:
+            print(f"Error processing log entry: {e}")

tests/__init__.py

@@ -0,0 +1,9 @@
+"""
+Test Suite for SecureX SDK
+Run all tests with: pytest tests/
+"""
+import sys
+import os
+
+# Add parent directory to path
+sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))

tests/test_coverage_client.py

@@ -0,0 +1,210 @@
+import pytest
+import asyncio
+import discord
+from securex import SecureX
+from unittest.mock import Mock, AsyncMock, patch, MagicMock
+
+@pytest.mark.asyncio
+class TestClientExhaustive:
+    """Exhaustive tests for securex/client.py to reach 100% coverage"""
+
+    async def test_audit_log_listener_success(self):
+        """Test on_audit_log_entry_create when everything is normal"""
+        bot = Mock(spec=discord.Client)
+        sx = SecureX(bot)
+        
+        # Mock queues
+        sx.action_queue = Mock(spec=asyncio.Queue)
+        sx.action_queue.put = AsyncMock()
+        sx.log_queue = Mock(spec=asyncio.Queue)
+        sx.log_queue.put = AsyncMock()
+        sx.cleanup_queue = Mock(spec=asyncio.Queue)
+        sx.cleanup_queue.put = AsyncMock()
+        sx.guild_queue = Mock(spec=asyncio.Queue)
+        sx.guild_queue.put = AsyncMock()
+        
+        # Get the listener
+        listeners = bot.event.call_args_list
+        listener_func = None
+        for call in listeners:
+            if call[0][0].__name__ == 'on_audit_log_entry_create':
+                listener_func = call[0][0]
+                break
+        
+        entry = Mock(spec=discord.AuditLogEntry)
+        await listener_func(entry)
+        
+        sx.action_queue.put.assert_called_once_with(entry)
+        sx.log_queue.put.assert_called_once_with(entry)
+        sx.cleanup_queue.put.assert_called_once_with(entry)
+        sx.guild_queue.put.assert_called_once_with(entry)
+
+    async def test_audit_log_listener_queue_full(self):
+        """Test on_audit_log_entry_create when queues are full"""
+        bot = Mock(spec=discord.Client)
+        sx = SecureX(bot)
+        
+        # Mock queues to raise QueueFull
+        sx.action_queue = Mock(spec=asyncio.Queue)
+        sx.action_queue.put = AsyncMock(side_effect=asyncio.QueueFull)
+        
+        # Get the listener
+        listeners = bot.event.call_args_list
+        listener_func = None
+        for call in listeners:
+            if call[0][0].__name__ == 'on_audit_log_entry_create':
+                listener_func = call[0][0]
+                break
+        
+        assert listener_func is not None
+        
+        entry = Mock(spec=discord.AuditLogEntry)
+        # Should catch QueueFull and print warning
+        await listener_func(entry)
+
+    async def test_registered_events(self):
+        """Test all registered Discord event listeners in client.py"""
+        bot = Mock(spec=discord.Client)
+        sx = SecureX(bot)
+        
+        event_mappers = {}
+        for call in bot.event.call_args_list:
+            func = call[0][0]
+            event_mappers[func.__name__] = func
+            
+        # 1. on_guild_channel_delete
+        channel = Mock()
+        sx.channel_handler.on_channel_delete = AsyncMock()
+        await event_mappers['on_guild_channel_delete'](channel)
+        sx.channel_handler.on_channel_delete.assert_called_once_with(channel)
+
+        # 2. on_guild_channel_update
+        before, after = Mock(), Mock()
+        sx.channel_handler.on_channel_update = AsyncMock()
+        await event_mappers['on_guild_channel_update'](before, after)
+        sx.channel_handler.on_channel_update.assert_called_once_with(before, after)
+
+        # 3. on_guild_role_delete
+        role = Mock()
+        sx.role_handler.on_role_delete = AsyncMock()
+        await event_mappers['on_guild_role_delete'](role)
+        sx.role_handler.on_role_delete.assert_called_once_with(role)
+
+        # 4. on_guild_role_update
+        before_r, after_r = Mock(), Mock()
+        sx.role_handler.on_role_update = AsyncMock()
+        await event_mappers['on_guild_role_update'](before_r, after_r)
+        sx.role_handler.on_role_update.assert_called_once_with(before_r, after_r)
+
+        # 5. on_member_ban
+        guild, user = Mock(), Mock()
+        sx.member_handler.on_member_ban = AsyncMock()
+        await event_mappers['on_member_ban'](guild, user)
+        sx.member_handler.on_member_ban.assert_called_once_with(guild, user)
+
+        # 6. on_member_update
+        before_m, after_m = Mock(), Mock()
+        sx.member_handler.on_member_update = AsyncMock()
+        await event_mappers['on_member_update'](before_m, after_m)
+        sx.member_handler.on_member_update.assert_called_once_with(before_m, after_m)
+
+    async def test_callback_decorators_and_triggering(self):
+        """Test @sx.on decorators and _trigger_callbacks logic"""
+        bot = Mock()
+        sx = SecureX(bot)
+        
+        # Test async callback
+        mock_async_cb = AsyncMock()
+        sx.on('threat_detected')(mock_async_cb)
+        
+        # Test sync callback
+        mock_sync_cb = MagicMock()
+        sx.on('threat_detected')(mock_sync_cb)
+        
+        await sx._trigger_callbacks('threat_detected', "arg1", extra="data")
+        
+        mock_async_cb.assert_called_once_with("arg1", extra="data")
+        mock_sync_cb.assert_called_once_with("arg1", extra="data")
+        
+        # Test callback with exception
+        fail_cb = Mock(side_effect=Exception("Callback failed"))
+        sx.on('threat_detected')(fail_cb)
+        # Should not raise exception
+        await sx._trigger_callbacks('threat_detected')
+        
+        # Test triggering non-existent event
+        await sx._trigger_callbacks('non_existent')
+        
+        # Test convenience properties
+        sx.on_threat_detected(lambda x: x)
+        sx.on_backup_completed(lambda x: x)
+        sx.on_restore_completed(lambda x: x)
+        sx.on_whitelist_changed(lambda x: x)
+        
+        assert len(sx._callbacks['threat_detected']) == 4 # 2 from before + 1 fail + 1 lambda
+        assert len(sx._callbacks['backup_completed']) == 1
+
+    async def test_enable_full(self):
+        """Test sx.enable with full set of parameters for max coverage"""
+        bot = Mock()
+        sx = SecureX(bot)
+        
+        sx.whitelist.preload_all = AsyncMock()
+        sx.whitelist.get_all = AsyncMock(return_value=[1, 2])
+        sx.whitelist.add = AsyncMock()
+        sx.backup_manager.preload_all = AsyncMock()
+        sx.backup_manager.create_backup = AsyncMock(return_value=Mock(success=True, channel_count=5, role_count=10))
+        
+        # Mock start_auto_backup and start_auto_refresh manually
+        sx.backup_manager.start_auto_backup = MagicMock()
+        sx.backup_manager.start_auto_refresh = MagicMock()
+        
+        await sx.enable(
+            guild_id=123,
+            whitelist=[3],
+            auto_backup=True,
+            punishments={"channel_delete": "kick"},
+            notify_user=True
+        )
+        
+        assert sx.whitelist_cache[123] == {1, 2, 3}
+        assert sx.punishment_cache[123]["channel_delete"] == "kick"
+        sx.backup_manager.start_auto_backup.assert_called_once()
+        sx.backup_manager.start_auto_refresh.assert_called_once()
+
+    async def test_enable_minimal(self):
+        """Test sx.enable with minimal parameters"""
+        bot = Mock()
+        sx = SecureX(bot)
+        
+        # Mocking to avoid real side effects
+        sx.whitelist.preload_all = AsyncMock()
+        sx.backup_manager.preload_all = AsyncMock()
+        sx.backup_manager.start_auto_backup = MagicMock()
+        sx.backup_manager.start_auto_refresh = MagicMock()
+        
+        await sx.enable(auto_backup=False)
+        
+        # These should NOT be called if auto_backup is False
+        sx.backup_manager.start_auto_backup.assert_not_called()
+        sx.backup_manager.start_auto_refresh.assert_not_called()
+
+    async def test_proxy_methods(self):
+        """Test create_backup, restore_channel, and restore_role proxy methods"""
+        bot = Mock()
+        sx = SecureX(bot)
+        
+        # restore_channel
+        sx.backup_manager.restore_channel = AsyncMock(return_value="new_channel")
+        res = await sx.restore_channel(Mock(), Mock())
+        assert res == "new_channel"
+        
+        # restore_role
+        sx.backup_manager.restore_role = AsyncMock(return_value=True)
+        res = await sx.restore_role(Mock(), 999)
+        assert res is True
+        
+        # create_backup
+        sx.backup_manager.create_backup = AsyncMock(return_value="backup_info")
+        res = await sx.create_backup(123)
+        assert res == "backup_info"