bot.py

import os
import asyncio
import threading
import signal
import time
import platform
import socket
import getpass
from pathlib import Path
import psutil
import discord
from discord.ext import commands
from dotenv import load_dotenv

load_dotenv()

TOKEN = os.getenv("DISCORD_TOKEN")
if not TOKEN:
    raise RuntimeError("DISCORD_TOKEN environment variable not set")
CHANNEL_ID = os.getenv("CHANNEL_ID")
if not CHANNEL_ID:
    raise RuntimeError("CHANNEL_ID environment variable not set")

intents = discord.Intents.default()
intents.message_content = True

client = commands.Bot(command_prefix="!", intents=intents)
client.remove_command('help')

@client.event
async def on_ready():
    print(f"Bot logged in as {client.user}")

@client.command(name="sysinfo")
async def sysinfo(ctx):
    """
    Grabs the PC current system information
    """
    info = get_system_info()
    message = "**🖥️ System Information**\n\n"
    for key, value in info.items():
        message += f"**{key}:** {value}\n"
    await ctx.send(message)

@client.command(name="pwd")
async def pwd(ctx):
    """
    Prints currnet working directory
    """
    await ctx.send(f"📁 **Current working directory:**\n```\n{os.getcwd()}\n```")

@client.command(name="ls")
async def ls(ctx):
    """
    List the items in the current directory
    """
    cwd = Path.cwd()

    dirs = []
    files = []

    for item in cwd.iterdir():
        if item.is_dir():
            dirs.append(f"[DIR]  {item.name}")
        else:
            files.append(f"       {item.name}")

    output = "\n".join(sorted(dirs) + sorted(files)) or "(empty)"

    await ctx.send(f"📂 **Listing: {cwd}**\n```\n{output}\n```")

@client.command(name="cd")
async def cd(ctx, *, path: str):
    """
    Change directory
    e.g !cd /home

    use !cd ../ to go back!
    """
    new_path = (Path.cwd() / path).resolve()

    if not new_path.exists():
        await ctx.send(f"❌ Path does not exist:\n```\n{new_path}\n```")
        return
    if not new_path.is_dir():
        await ctx.send(f"❌ Not a directory:\n```\n{new_path}\n```")
        return

    os.chdir(new_path)
    await ctx.send(f"📁 **Directory changed to:**\n```\n{new_path}\n```")

@client.command(name="whoami")
async def whoami(ctx):
    """
    Prints the current user login
    """
    user = getpass.getuser()
    await ctx.send(f"👤 **Current user:**\n```\n{user}\n```")

@client.command(name="search")
async def search(ctx, *, query: str):
    """
    Searches for the item
    e.g: !search hi.txt
    """
    root = Path.cwd()
    pattern = query.lower()
    max_depth = 5   

    results = []
    try:
        fast_search(root, pattern, max_depth, results)
    except Exception:
        pass

    if not results:
        await ctx.send(f"🔍 **No results for:** `{query}`")
        return

    output = "🔍 **Search results:**\n```\n" + "\n".join(results) + "\n```"

    if len(output) > 1900:
        chunks = [output[i:i+1900] for i in range(0, len(output), 1900)]
        for part in chunks:
            await ctx.send(part)
    else:
        await ctx.send(output)

@client.command(name="export_here")
async def export_here(ctx):
    """
    Export all the files in the current directory
    """
    cwd = Path.cwd()

    files = [f for f in cwd.iterdir() if f.is_file()]

    if not files:
        await ctx.send("📂 **No files to export in this directory.**")
        return

    preview = "\n".join(f"- {f.name}" for f in files)
    await ctx.send(
        f"📦 **The following files in `{cwd}` will be exported:**\n"
        f"```\n{preview}\n```\n"
        f"Type **yes** to confirm."
    )

    def check(msg):
        return msg.author == ctx.author and msg.channel == ctx.channel

    try:
        reply = await client.wait_for("message", timeout=20, check=check)
    except asyncio.TimeoutError:
        await ctx.send("⌛ **Timed out. Export cancelled.**")
        return

    if reply.content.lower().strip() != "yes":
        await ctx.send("❌ **Export cancelled.**")
        return

    zip_path = cwd / "export_here.zip"
    try:
        import zipfile
        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
            for f in files:
                z.write(f, arcname=f.name)

        await ctx.send("📤 **Exporting files...**")
        await ctx.send(file=discord.File(zip_path))

    except Exception as e:
        await ctx.send(f"❌ Error: `{e}`")
    finally:
        if zip_path.exists():
            zip_path.unlink()

    await ctx.send("✅ **Export complete.**")

@client.command(name="help")
async def help_command(ctx):
    """Display all available commands"""
    help_message = """
```ansi
Available Commands:

System Information:
  !sysinfo          - Display detailed system information (hostname, IP, OS, CPU, RAM, disk)
  !whoami           - Show current user login

File Navigation:
  !pwd              - Show current working directory
  !ls               - List files and directories in the current location
  !cd <path>        - Change directory (use !cd ../ to go back)
  !tree [depth] [-a] - Show directory tree structure
                      Examples: !tree, !tree 5, !tree -a, !tree 4 -a

File Operations:
  !search <query>   - Search for files/folders by name in current directory
                      Example: !search hi.txt
  !export_here      - Export all files in the current directory as a zip file

Other:
  !help             - Show this help message

Tips:
- All commands are case-sensitive and must start with !
- Use !cd ../ to navigate to parent directory
- Use !tree -a to show hidden files
```
    """
    await ctx.send(help_message)


def generate_tree(path: Path, prefix: str = "", depth: int = 0, max_depth: int = 3, show_hidden: bool = False):
    """Recursive directory tree generator with clean formatting."""
    if depth > max_depth:
        return prefix + "└── ... (depth limit reached)\n"

    tree_output = ""

    items = [
        item for item in path.iterdir()
        if show_hidden or not item.name.startswith(".")
    ]

    items = sorted(items, key=lambda x: (not x.is_dir(), x.name.lower()))
    count = len(items)

    for idx, item in enumerate(items):
        connector = "└── " if idx == count - 1 else "├── "
        tree_output += prefix + connector + item.name + "\n"

        if item.is_dir():
            extension = "    " if idx == count - 1 else "│   "
            tree_output += generate_tree(
                item,
                prefix + extension,
                depth + 1,
                max_depth,
                show_hidden
            )

    return tree_output


@client.command(name="tree")
async def tree(ctx, *args):
    """
    Prints a tree of the current directory
      !tree               → hidden ignored, depth=3
      !tree 4             → depth=4
      !tree -a            → show hidden
      !tree 4 -a          → depth=4 + hidden
    """
    depth = 3
    show_hidden = False

    for arg in args:
        if arg.isdigit():
            depth = int(arg)
        if arg in ("-a", "--all", "a"):
            show_hidden = True

    root = Path.cwd()

    output = f"📁 **Tree for: {root}**\n```\n"
    output += generate_tree(root, max_depth=depth, show_hidden=show_hidden)
    output += "```"

    # Discord message size limit
    if len(output) > 1800:
        chunks = [output[i:i+1800] for i in range(0, len(output), 1800)]
        for part in chunks:
            await ctx.send(part)
    else:
        await ctx.send(output)

def fast_search(root: Path, pattern: str, depth: int, results: list):
    if depth < 0:
        return

    for item in root.iterdir():
        name = item.name.lower()

        if pattern in name:
            results.append(str(item))

        if item.is_dir():
            try:
                fast_search(item, pattern, depth - 1, results)
            except Exception:
                continue

def get_system_info():
    info = {}
    info["Hostname"] = socket.gethostname()
    info["IP Address"] = socket.gethostbyname(info["Hostname"])
    info["OS"] = platform.platform()
    info["Architecture"] = platform.machine()
    info["Processor"] = platform.processor()
    info["Physical Cores"] = psutil.cpu_count(logical=False)
    info["Total Cores"] = psutil.cpu_count(logical=True)

    ram = psutil.virtual_memory()
    info["Total RAM"] = f"{ram.total / (1024**3):.2f} GB"
    info["Available RAM"] = f"{ram.available / (1024**3):.2f} GB"

    disk = psutil.disk_usage('/')
    info["Disk Total"] = f"{disk.total / (1024**3):.2f} GB"
    info["Disk Free"] = f"{disk.free / (1024**3):.2f} GB"

    return info

shutdown_flag = False

def handle_exit(signum, frame):
    global shutdown_flag
    shutdown_flag = True
    print("\nShutting down cleanly...")

signal.signal(signal.SIGINT, handle_exit)
signal.signal(signal.SIGTERM, handle_exit)

def run_bot():
    asyncio.run(client.start(TOKEN))

threading.Thread(target=run_bot, daemon=True).start()

while not shutdown_flag:
    try:
        time.sleep(0.5)  
    except KeyboardInterrupt:
        break

try:
    asyncio.run(client.close())
except Exception:
    pass

print("Bot stopped.")