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 .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")

# Ensure all required env vars are present
for var in (
    "DISCORD_TOKEN", "GUILD_ID", "ROLE_ID",
    "LNBITS_URL", "LNBITS_API_KEY", "PRICE",
    "CHANNEL_ID", "DATABASE_URL"
):
    if not globals()[var]:
        logger.critical(f"Environment variable {var} is missing.")
        exit(1)

# Build WS URL from LNBITS_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}")
    data = pending_invoices.pop(payment_hash, None)
    if data is None:
        logger.info(f"No pending invoice for hash={payment_hash}")
        return
    user_id, interaction = data

    # Convert from millisatoshis to sats
    raw_msat = payment_obj.get("amount", 0)
    sat_amount = raw_msat // 1000

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

    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)
    channel = bot.get_channel(CHANNEL_ID)

    if not all([member, role, channel]):
        logger.error("Member, role, or channel not found")
        return

    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
            )
            await channel.send(
                f"🎉 {member.mention} has paid **{sat_amount} sats** "
                f"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:
        await channel.send(
            f"ℹ️ {member.mention} paid again (hash={payment_hash}), "
            f"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)
                        if hsh and pay.get("status") == "success" and amt > 0:
                            logger.info(f"💡 Payment received hash={hsh}, amt={amt}")
                            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})")
    try:
        synced = await bot.tree.sync()
        logger.info(f"✅ Synced {len(synced)} command(s).")
    except Exception:
        logger.exception("Failed to sync commands")
    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

    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"
    }

    # Create invoice
    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()

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

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

    # Generate QR code
    def make_qr_buffer():
        img = qrcode.make(payment_request)
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        buf.seek(0)
        return buf

    qr_buffer = await loop.run_in_executor(None, make_qr_buffer)
    qr_file   = File(qr_buffer, filename="invoice.png")

    embed = Embed(
        title="⚡ Please Pay ⚡",
        description=INVOICE_MESSAGE_TEMPLATE,
        color=discord.Color.gold()
    )
    embed.add_field(name="Invoice", value=f"```{payment_request}```", 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: {payment_hash}")

    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)