Stable Diffusion Parasitic Bot

bot.py

#!/usr/bin/env python
# coding: utf-8

import os
import tempfile
import time

import requests
from discord import File
from discord.ext import commands
from dotenv import load_dotenv
from loguru import logger
from selenium import webdriver
from selenium.common.exceptions import StaleElementReferenceException, NoSuchElementException  # noqa
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys

load_dotenv()

CHANNEL_LINK = os.environ['CHANNEL_LINK']
CHROME_BIN = os.environ['CHROME_BIN']
CHROMEDRIVER_EXEC = os.environ['CHROMEDRIVER_EXEC']
BOT_TOKEN = os.environ['BOT_TOKEN']
BOT_CHANNEL_ID = int(os.environ['BOT_CHANNEL_ID'])
HELPER_USERNAME = os.environ['HELPER_USERNAME']
HELPER_TOKEN = os.environ['HELPER_TOKEN']
PROCESSED = []

bot = commands.Bot(command_prefix='!')


def setup_driver():
    options = webdriver.ChromeOptions()
    options.add_argument('--headless')
    options.add_argument('--mute-audio')
    options.add_argument('--start-maximized')
    options.add_argument('--no-sandbox')
    options.add_argument('--disable-dev-shm-usage')
    options.binary_location = CHROME_BIN

    service = Service(CHROMEDRIVER_EXEC)
    driver = webdriver.Chrome(options=options, service=service)

    driver.get('https://discord.com/login')
    driver.execute_script("let token =\"" + HELPER_TOKEN + """\";
                function login(token) {
                setInterval(() => {
                      document.body.appendChild(
                      document.createElement `iframe`
                      ).contentWindow.localStorage.token = `"${token}"`
                    }, 50);
                    setTimeout(() => {
                      location.reload();
                    }, 2500);
                 }
                login(token);""")

    time.sleep(8)
    driver.get(CHANNEL_LINK)
    time.sleep(1)
    elms = driver.find_elements(By.TAG_NAME, 'li')
    for e in elms:
        try:
            if (HELPER_USERNAME in e.text and 'DreamBotMothership'
                    in e.text) and e.text not in PROCESSED:
                PROCESSED.append(e.text)
        except (NoSuchElementException, StaleElementReferenceException) as e:
            logger.warning(e)
    logger.info('Ready...')
    return driver


def lock_channel(ctx, send_messages=False):
    overwrite = ctx.channel.overwrites_for(ctx.guild.default_role)
    overwrite.send_messages = send_messages
    return ctx.channel.set_permissions(ctx.guild.default_role,
                                       overwrite=overwrite)


@bot.command(name='dream')
@commands.max_concurrency(1, wait=True)
@commands.cooldown(1, 60)
async def dream(ctx, *args):
    if ctx.channel.id != BOT_CHANNEL_ID:
        logger.warning('Called in the wrong channel!')
        return

    await lock_channel(ctx)

    prompt = ' '.join(args).strip()
    message = f'!dream {prompt} -g'
    ActionChains(driver).send_keys(message).send_keys(Keys.ENTER).perform()

    BREAK_PARENT = False
    n = 0
    while True:
        n += 1
        elms = driver.find_elements(By.TAG_NAME, 'li')
        logger.debug(f'Waiting... ({n})')
        for e in elms:
            try:
                if (HELPER_USERNAME in e.text and 'DreamBotMothership'
                        in e.text) and e.text not in PROCESSED:
                    logger.debug(e.text)
                    try:
                        output_url = e.find_element(By.TAG_NAME,
                                                    'a').get_attribute('href')
                    except NoSuchElementException:
                        logger.warning('NoSuchElementException')
                        continue
                    with tempfile.NamedTemporaryFile(suffix='.png') as f:
                        r = requests.get(output_url)
                        f.write(r.content)
                        f.seek(0)
                        output = File(f.name)
                        await lock_channel(ctx, send_messages=True)
                        await ctx.send(ctx.message.author.mention, file=output)
                        await lock_channel(ctx)
                    PROCESSED.append(e.text)
                    BREAK_PARENT = True
                    break
            except StaleElementReferenceException:
                logger.warning('StaleElementReferenceException')
                elms = driver.find_elements(By.TAG_NAME, 'li')
        if BREAK_PARENT:
            logger.info('Breaking...')
            break
        else:
            time.sleep(2)

    await lock_channel(ctx, send_messages=True)


driver = setup_driver()
bot.run(BOT_TOKEN)

requirements.txt

requests>=2.28.1
discord.py>=1.7.3
python-dotenv>=0.20.0
loguru>=0.6.0
selenium>=4.3.0