discord-lnbits-bot/discord_lnbits_bot.py

#!/usr/bin/env python3
import os
import asyncio
import json
import io
import qrcode
import requests
import logging
from datetime import datetime
from dotenv import load_dotenv

import discord
from discord import File, Embed
from discord.ext import commands

import websockets

from sqlalchemy import (
    create_engine, Column, String, Integer, DateTime
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

# ─── Environment & Logging ────────────────────────────────────────────────────
load_dotenv()  # loads variables from .env into os.environ

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)7s %(name)s | %(message)s"
)
logger = logging.getLogger("discord_lnbits_bot")

# ─── Config from ENV ──────────────────────────────────────────────────────────
DISCORD_TOKEN            = os.getenv("DISCORD_TOKEN")
GUILD_ID                 = int(os.getenv("GUILD_ID", 0))
ROLE_ID                  = int(os.getenv("ROLE_ID", 0))
LNBITS_URL               = os.getenv("LNBITS_URL", "").rstrip("/")
LNBITS_API_KEY           = os.getenv("LNBITS_API_KEY", "")
PRICE                    = int(os.getenv("PRICE", 0))
CHANNEL_ID               = int(os.getenv("CHANNEL_ID", 0))
INVOICE_MESSAGE_TEMPLATE = os.getenv("INVOICE_MESSAGE", "Invoice for your purchase.")
COMMAND_NAME             = os.getenv("COMMAND_NAME", "support")
DATABASE_URL             = os.getenv("DATABASE_URL")  # e.g. postgres://user:pass@db:5432/bot

# Sanity checks
for var_name in (
    "DISCORD_TOKEN", "GUILD_ID", "ROLE_ID",
    "LNBITS_URL", "LNBITS_API_KEY", "PRICE", 
    "CHANNEL_ID", "DATABASE_URL"
):
    if not globals()[var_name]:
        logger.critical(f"Environment variable {var_name} is missing.")
        exit(1)

# Construct WS URL
if LNBITS_URL.startswith("https://"):
    base_ws = LNBITS_URL.replace("https://", "wss://", 1)
elif LNBITS_URL.startswith("http://"):
    base_ws = LNBITS_URL.replace("http://", "ws://", 1)
else:
    logger.critical("LNBITS_URL must start with http:// or https://")
    exit(1)

LNBITS_WS_URL = f"{base_ws}/api/v1/ws/{LNBITS_API_KEY}"


# ─── Database Setup ───────────────────────────────────────────────────────────
Base = declarative_base()

class Payment(Base):
    __tablename__ = "payments"
    id            = Column(Integer, primary_key=True, autoincrement=True)
    payment_hash  = Column(String(256), unique=True, nullable=False, index=True)
    user_id       = Column(String(64), nullable=False, index=True)
    amount        = Column(Integer, nullable=False)
    paid_at       = Column(DateTime, default=datetime.utcnow, nullable=False)

engine = create_engine(DATABASE_URL, future=True)
SessionLocal = sessionmaker(bind=engine, expire_on_commit=False)
Base.metadata.create_all(engine)
logger.info("✅ Database tables ensured.")


# ─── Discord Bot Setup ────────────────────────────────────────────────────────
intents = discord.Intents.default()
intents.members = True

bot = commands.Bot(command_prefix="!", intents=intents)
pending_invoices = {}  # payment_hash -> (user_id, interaction)

# ─── Role Assignment Handler ──────────────────────────────────────────────────
async def assign_role_after_payment(payment_hash: str, payment_obj: dict):
    logger.debug(f"ENTER assign_role_after_payment for hash={payment_hash}")
    user_id, interaction = pending_invoices.pop(payment_hash, (None, None))
    if user_id is None:
        logger.info(f"No pending invoice for hash={payment_hash}")
        return

    # Record to database
    amt = payment_obj.get("amount", 0)
    session = SessionLocal()
    try:
        payment = Payment(
            payment_hash=payment_hash,
            user_id=str(user_id),
            amount=amt,
            paid_at=datetime.utcnow()
        )
        session.add(payment)
        session.commit()
        logger.info(f"💾 Logged payment {payment_hash} for user {user_id} ({amt} sats).")
    except Exception:
        session.rollback()
        logger.exception("DB error saving payment")
    finally:
        session.close()

    # Fetch guild, member, role
    guild = bot.get_guild(GUILD_ID)
    if not guild:
        logger.error(f"Guild {GUILD_ID} not found")
        return

    member = guild.get_member(user_id) or await guild.fetch_member(user_id)
    role = guild.get_role(ROLE_ID)
    if not role:
        logger.error(f"Role {ROLE_ID} not found in guild {guild.id}")
        return

    # Assign role if needed
    if role not in member.roles:
        logger.debug(f"Adding role '{role.name}' to {member.display_name}")
        try:
            await asyncio.wait_for(
                member.add_roles(role, reason="Paid via Lightning"),
                timeout=10
            )
            channel = bot.get_channel(CHANNEL_ID)
            await channel.send(
                f"🎉 {member.mention} has paid **{amt} sats** and received the **{role.name}** role!"
            )
            logger.info(f"✅ Role '{role.name}' assigned to {member.display_name}")
        except Exception:
            logger.exception("Error assigning role or sending confirmation")
    else:
        # Already has role
        channel = bot.get_channel(CHANNEL_ID)
        await channel.send(
            f"ℹ️ {member.mention} paid again (hash={payment_hash}), role **{role.name}** already assigned."
        )
        logger.info(f"User {member.display_name} already had role; notified channel.")


# ─── LNbits WebSocket Listener ────────────────────────────────────────────────
async def lnbits_ws_listener():
    await bot.wait_until_ready()
    logger.info(f"👂 Connecting to LNbits WS at {LNBITS_WS_URL}")
    while True:
        try:
            async with websockets.connect(LNBITS_WS_URL) as ws:
                logger.info("✅ WS connected")
                async for msg in ws:
                    logger.debug(f"WS ► {msg}")
                    try:
                        data = json.loads(msg)
                        pay = data.get("payment", {})
                        hsh = pay.get("checking_id") or pay.get("payment_hash")
                        amt = pay.get("amount", 0)
                        # LNbits uses status="success" on paid
                        if hsh and pay.get("status") == "success" and amt > 0:
                            logger.info(
                                f"💡 Payment received hash={hsh}, amt={amt}"
                            )
                            # schedule role assignment
                            bot.loop.create_task(
                                assign_role_after_payment(hsh, pay)
                            )
                        else:
                            logger.debug("Ignored WS update (not a paid invoice)")
                    except json.JSONDecodeError:
                        logger.warning("WS ► Non-JSON payload")
        except Exception:
            logger.exception("WS listener error; reconnecting in 10s")
            await asyncio.sleep(10)


# ─── Slash Command Registration ───────────────────────────────────────────────
@bot.event
async def on_ready():
    logger.info(f"Logged in as {bot.user} (ID: {bot.user.id})")
    # sync commands
    try:
        synced = await bot.tree.sync()
        logger.info(f"✅ Synced {len(synced)} command(s).")
    except Exception:
        logger.exception("Failed to sync commands")

    # start WS listener
    bot.loop.create_task(lnbits_ws_listener())


@bot.tree.command(name=COMMAND_NAME, description="Pay to get your role via Lightning")
async def dynamic_command(interaction: discord.Interaction):
    invoice_ch = bot.get_channel(CHANNEL_ID)
    if invoice_ch is None:
        await interaction.response.send_message(
            "❌ Invoice channel not found. Check your config.",
            ephemeral=True
        )
        return

    # offload blocking REST call + QR gen
    loop = asyncio.get_running_loop()
    invoice_data = {"out": False, "amount": PRICE, "memo": f"Role for {interaction.user.display_name}"}
    headers = {"X-Api-Key": LNBITS_API_KEY, "Content-Type": "application/json"}

    def create_invoice():
        resp = requests.post(
            f"{LNBITS_URL}/api/v1/payments",
            json=invoice_data,
            headers=headers,
            timeout=10
        )
        resp.raise_for_status()
        return resp.json()

    # Create invoice
    try:
        invoice_json = await loop.run_in_executor(None, create_invoice)
    except Exception as e:
        logger.exception("Error creating LNbits invoice")
        await interaction.response.send_message(
            "❌ Failed to generate invoice. Try again later.",
            ephemeral=True
        )
        return

    pay_req  = invoice_json["bolt11"]
    pay_hash = invoice_json["payment_hash"]
    pending_invoices[pay_hash] = (interaction.user.id, interaction)
    logger.info(f"Stored pending invoice {pay_hash} for user {interaction.user.id}")

    # Generate QR
    qr_bytes = await loop.run_in_executor(None, lambda: qrcode.make(pay_req).get_image().tobytes())
    qr_file  = File(io.BytesIO(qr_bytes), filename="invoice.png")

    embed = Embed(
        title="⚡ Please Pay ⚡",
        description=INVOICE_MESSAGE_TEMPLATE,
        color=discord.Color.gold()
    )
    embed.add_field(name="Invoice", value=f"```{pay_req}```", inline=False)
    embed.add_field(name="Amount", value=f"{PRICE} sats", inline=True)
    embed.set_image(url="attachment://invoice.png")
    embed.set_footer(text=f"Hash: {pay_hash}")

    # Send invoice & ack
    await invoice_ch.send(content=interaction.user.mention, embed=embed, file=qr_file)
    await interaction.response.send_message(
        f"✅ Invoice posted in {invoice_ch.mention}", ephemeral=True
    )


# ─── Run Bot ──────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    bot.run(DISCORD_TOKEN)