Dislocker/dislocker.py

import json
import discord
import os
import psycopg2
import hashlib
import string
import random

class Bot(discord.Client):
    def db_connect(self, host, db, port, user, password):
        self.db = psycopg2.connect(f"host={host} dbname={db} port={port} user={user} password={password}")

    def password_generate(self, length):
        numbers = string.digits  # (1)
        password = ''.join(random.choice(numbers) for _ in range(length))  # (2)
        return str(password)
    
    def hash_genarate(self, source):
        hashed = hashlib.md5(source.encode())
        return hashed.hexdigest()

    async def on_ready(self):
        print("ログイン成功")

    async def on_message(self, message):
        if message.author.bot:
            return
        
        if isinstance(message.channel, discord.DMChannel):
            user_id = message.author.id
            user_name = message.author.name
            msg_split = message.content.split()
            if msg_split[0] == "/password":
                if len(msg_split) <= 2:
                    await message.channel.send("PC番号、もしくはデバイス番号が入力されていません。")

                elif len(msg_split) == 3:
                    if msg_split[1].isdigit() and msg_split[2].isdigit():
                        if int(msg_split[1]) <= 10 and int(msg_split[1]) >= 1:
                            password = self.password_generate(4)
                            password_hash = self.hash_genarate(password)
                            cursor = self.db.cursor()
                            cursor.execute("SELECT * FROM club_member WHERE discord_userid = %s", (str(message.author.id),))
                            member_info = cursor.fetchall()
                            cursor.execute("SELECT * FROM pc_usage_history WHERE member_id=%s ORDER BY id DESC LIMIT 1", (member_info[0][0],))
                            history_info = cursor.fetchall()
                            if len(history_info) == 0:
                                cursor.execute("INSERT INTO pc_usage_history ( member_id, pc_number, device_number, start_use_time ) VALUES ( %s, %s, %s, current_timestamp)", (member_info[0][0], int(msg_split[1]), int(msg_split[2])))
                                self.db.commit()
                                cursor.execute("UPDATE pc_list SET using_user_id = %s WHERE pc_number = %s", (member_info[0][0], msg_split[1]))
                                self.db.commit()
                                cursor.execute("UPDATE pc_list SET password_hash = %s WHERE pc_number = %s", (password_hash, msg_split[1]))
                                self.db.commit()
                                
                                await message.channel.send(f"使用が開始されました。\nパスワード | {password}\nPC番号 | {msg_split[1]}\nデバイス番号 | {msg_split[2]}\n")
                                await self.get_channel(bot_config["log_channel_id"]).send(f'{member_info[0][2]}がPC{msg_split[1]}を使用しています')
                                
                            elif not history_info[0][5]:
                                await message.channel.send(f"もう使用されているようです。解除するには /stop で使用終了をお知らせください。\n使用中のPC番号 | {str(history_info[0][2])}\n使用中のデバイス番号 | {str(history_info[0][3])}\n使用開始時刻 | {str(history_info[0][4])}\n使用目的 | {str(history_info[0][6])}")
                            else:
                                cursor.execute("INSERT INTO pc_usage_history ( member_id, pc_number, device_number, start_use_time ) VALUES ( %s, %s, %s, current_timestamp)", (member_info[0][0], int(msg_split[1]), int(msg_split[2])))
                                self.db.commit()
                                cursor.execute("UPDATE pc_list SET using_user_id = %s WHERE pc_number = %s", (member_info[0][0], msg_split[1]))
                                self.db.commit()
                                cursor.execute("UPDATE pc_list SET password_hash = %s WHERE pc_number = %s", (password_hash, msg_split[1]))
                                self.db.commit()
                                
                                await message.channel.send(f"使用が開始されました。\nパスワード | {password}\nPC番号 | {msg_split[1]}\nデバイス番号 | {msg_split[2]}\n")
                                await self.get_channel(bot_config["log_channel_id"]).send(f'{member_info[0][2]}がPC{msg_split[1]}を使用しています')
                        else:
                            await message.channel.send("パソコンの台数が\n# だめです")
                    else:
                        await message.channel.send("構文が不正です。")

                elif len(msg_split) >= 4:
                    if msg_split[1].isdigit() and msg_split[2].isdigit():
                        if int(msg_split[1]) <= 10 and int(msg_split[1]) >= 1:
                            password = self.password_generate(4)
                            password_hash = self.hash_genarate(password)
                            cursor = self.db.cursor()
                            cursor.execute("SELECT * FROM club_member WHERE discord_userid = %s", (str(message.author.id),))
                            member_info = cursor.fetchall()
                            cursor.execute("SELECT * FROM pc_usage_history WHERE member_id=%s ORDER BY id DESC LIMIT 1", (member_info[0][0],))
                            history_info = cursor.fetchall()
                            if len(history_info) == 0:
                                cursor.execute("INSERT INTO pc_usage_history ( member_id, pc_number, device_number, start_use_time ) VALUES ( %s, %s, %s, current_timestamp)", (member_info[0][0], int(msg_split[1]), int(msg_split[2])))
                                self.db.commit()
                                cursor.execute("UPDATE pc_list SET using_user_id = %s WHERE pc_number = %s", (member_info[0][0], msg_split[1]))
                                self.db.commit()
                                cursor.execute("UPDATE pc_list SET password_hash = %s WHERE pc_number = %s", (password_hash, msg_split[1]))
                                self.db.commit()
                                
                                await message.channel.send(f"使用が開始されました。\nパスワード | {password}\nPC番号 | {msg_split[1]}\nデバイス番号 | {msg_split[2]}\n")
                                await self.get_channel(bot_config["log_channel_id"]).send(f'{member_info[0][2]}がPC{msg_split[1]}を使用しています')
                            elif not history_info[0][5]:
                                await message.channel.send(f"もう使用されているようです。解除するには /stop で使用終了をお知らせください。\n使用中のPC番号 | {str(history_info[0][2])}\n使用中のデバイス番号 | {str(history_info[0][3])}\n使用開始時刻 | {str(history_info[0][4])}\n使用目的 | {str(history_info[0][6])}")
                            else:
                                cursor.execute("INSERT INTO pc_usage_history ( member_id, pc_number, device_number, start_use_time, use_detail ) VALUES ( %s, %s, %s, current_timestamp, %s)", (member_info[0][0], int(msg_split[1]), int(msg_split[2]), str(msg_split[3])))
                                self.db.commit()
                                cursor.execute("UPDATE pc_list SET using_user_id = %s WHERE pc_number = %s", (member_info[0][0], msg_split[1]))
                                self.db.commit()
                                cursor.execute("UPDATE pc_list SET password_hash = %s WHERE pc_number = %s", (password_hash, msg_split[1]))
                                self.db.commit()
                                await message.channel.send(f"使用が開始されました。\nパスワード | {password}\nPC番号 | {msg_split[1]}\nデバイス番号 | {msg_split[2]}\n")
                                await self.get_channel(bot_config["log_channel_id"]).send(f'{member_info[0][2]}がPC{msg_split[1]}を使用しています')
                        else:
                            await message.channel.send("パソコンの台数が\n# だめです")
                    else:
                        await message.channel.send("構文が不正です。")
                
            
            elif msg_split[0] == "/stop":
                cursor = self.db.cursor()
                cursor.execute("SELECT * FROM club_member WHERE discord_userid = %s", (str(message.author.id),))
                member_info = cursor.fetchall()
                cursor.execute("SELECT * FROM pc_usage_history WHERE member_id=%s ORDER BY id DESC LIMIT 1", (member_info[0][0],))
                history_info = cursor.fetchall()
                if history_info[0][5]:
                    await message.channel.send("使用されていないようです...")
                else:
                    cursor.execute("UPDATE pc_usage_history SET end_use_time = current_timestamp WHERE id = %s", (history_info[0][0],))
                    self.db.commit()
                    cursor.execute("UPDATE pc_list SET using_user_id = NULL WHERE pc_number = %s", (history_info[0][2],))
                    self.db.commit()
                    cursor.execute("UPDATE pc_list SET password_hash = NULL WHERE pc_number = %s", (history_info[0][2],))
                    self.db.commit()
                    await message.channel.send(f"使用が終了されました。")
                    await self.get_channel(bot_config["log_channel_id"]).send(f'{member_info[0][2]}がPC{history_info[0][2]}の使用を終了しました')

        else:

            if message.content == "/hello":
                await message.channel.send("こんにちは！！！！")

            else:
                await message.channel.send("どうしましたか？")


config_dir_path = "./config/"
db_config_path = config_dir_path + "db.json"
if not os.path.isfile(db_config_path):
    if not os.path.isdir(config_dir_path):
        os.mkdir(config_dir_path)

    db_config = {
        "host": "localhost",
        "db": "dislocker",
        "username": "user",
        "password": "example_pass",
        "port": "5432"
    }
    with open(db_config_path, "w") as w:
        json.dump(db_config, w, indent=4)
elif os.path.isfile(db_config_path):
    with open(db_config_path, "r") as r:
        db_config = json.load(r)

bot_config_path = config_dir_path + "bot.json"
if not os.path.isfile(bot_config_path):
    if not os.path.isdir(config_dir_path):
        os.mkdir(config_dir_path)

    bot_config = {
        "token": "TYPE HERE BOTS TOKEN KEY",
        "log_channel_id" : "TYPE HERE CHANNEL ID"
    }
    with open(bot_config_path, "w") as w:
        json.dump(bot_config, w, indent=4)

elif os.path.isfile(bot_config_path):
    with open(bot_config_path, "r") as r:
        bot_config = json.load(r)

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

bot = Bot(intents=intents)
bot.db_connect(db_config["host"], db_config["db"], db_config["port"], db_config["username"], db_config["password"])
bot.run(bot_config["token"])