up db test

test.py

@@ -1,17 +1,18 @@
 import json
-import os
+import logging
+import sqlite3
 import sys
 import threading
-import logging
-import discord
-import sqlite3
+from os import walk, path, makedirs, remove, rename
+from typing import List, Union, Any
 
-from os import walk
+import discord
 from discord.ext import commands
-from dislash import InteractionClient, Option, OptionType, SelectMenu, SelectOption
+from dislash import InteractionClient, Option, OptionType
 
 
 class DB:
+    @staticmethod
     def _prepare_db(guild):
         try:
             connect = sqlite3.connect('user.db')
@@ -24,6 +25,7 @@ class DB:
         except sqlite3.Error as error:
             logging.error("failed to connecct db", error)
 
+    @staticmethod
     def _work_with_db(db_func, data_turple):
         try:
             connect = sqlite3.connect('user.db')
@@ -35,18 +37,21 @@ class DB:
         except sqlite3.Error as error:
             logging.error("failed to connecct db", error)
 
-    def _fill_bd(name, userid, isbot, nick, guild):
+    @staticmethod
+    def _fill_bd(name: str, userid: int, isbot: bool, nick: str, guild: int):
         sqlite_insert_with_param = (f"""INSERT OR IGNORE INTO "{guild}"
                               (username, userid, nick, isbot) 
                               VALUES (?, ?, ?, ?)""")
         data_tuple = (name, userid, nick, isbot)
         DB._work_with_db(sqlite_insert_with_param, data_tuple)
 
+    @staticmethod
     def _add_audio(guildid, user, audio):
         sql_update_query = f"""UPDATE "{guildid}" set track = ? where userid = ?"""
         data_tuple = (audio, user)
         DB._work_with_db(sql_update_query, data_tuple)
 
+    @staticmethod
     def _read_db(guildid, user):
         try:
             sql_con = sqlite3.connect("user.db")
@@ -54,10 +59,11 @@ class DB:
             sql_read = f"""SELECT * FROM "{guildid}" where userid = {user}"""
             cursor.execute(sql_read)
             record = cursor.fetchone()
-            if record[4] is None:
-                return "None"
-            else:
-                return record[4]
+            # if record[4] is None:
+            #     return "None"
+            # else:
+            #     return record[4]
+            return record[4]
 
         except sqlite3.Error as error:
             logging.error("Failed to read sqlite table", error)
@@ -67,38 +73,76 @@ threading.current_thread().name = "main"
 logging.basicConfig(stream=sys.stdout, filemode='w', level='INFO',
                     format='%(asctime)s - %(levelname)s - %(threadName)s - %(message)s')
 
-with open("config.json", 'r') as f:
-    prefixes = json.load(f)
-default_prefix = "$"
+DEFAULT_PREFIX = "$"  # The prefix you want everyone to have if you don't define your own
+
+if not path.isfile('prefix.json'):
+    with open('prefix.json', 'w+') as f:
+        f.write("")
+    f.close()
 
 
-def prefix(bot, message):
-    id = message.guild.id
-    return prefixes.get(id, default_prefix)
+def determine_prefix(bot, msg):
+    guild = msg.guild
+    base = DEFAULT_PREFIX
+    try:
+        with open('prefix.json', 'r', encoding='utf-8') as fp:  # Open the JSON
+            jsonObject = json.load(fp)  # Load the custom prefixes
+
+    except:
+        jsonObject = {}
+
+    if guild:  # If the guild exists
+        try:
+            guild_conf = jsonObject[f"{guild.id}"]
+            try:
+                prefix = guild_conf["prefix"]  # Get the prefix
+                base = prefix
+            except:
+                pass
+        except KeyError:
+            pass
+
+    return base
 
 
 intents = discord.Intents.default()
 intents.members = True
-bot = commands.Bot(command_prefix=prefix, guild_subscriptions=True, intents=intents)
+bot = commands.Bot(command_prefix=determine_prefix, guild_subscriptions=True, intents=intents)
 inter_client = InteractionClient(bot)
 
 
-class Arg:
-    f = []
-    for filenames in walk('audio'):
-        f.extend(filenames)
-        break
-    f = f[2]
-    dict = {}
-    keys = range(len(f))
-    for i in keys:
+class Commands:
+    async def setprefix(ctx, prefix: str):
+        with open('prefix.json', 'r', encoding='utf-8') as fp:
+            try:
+                jsonObject = json.load(fp)
+            except:
+                jsonObject = {}
+            try:
+                jsonObject[f"{ctx.guild.id}"]["prefix"] = prefix
+            except KeyError:
+                guild = jsonObject[f"{ctx.guild.id}"] = {'prefix': ''}
+                new = {"prefix": prefix}
+                guild.update(new)
+                jsonObject[str(ctx.guild.id)] = guild
+
+        await ctx.send(f"Prefix set to: `{prefix}`")
+
+        with open("prefix.json", "w") as fl:
+            json.dump(jsonObject, fl)
+            fl.close()
+
+    def list_files(folder):
+        f = []
+        for filenames in walk(folder):
+            f.extend(filenames)
+            break
+        f = f[2]
+        files = {}
         for x in f:
-            dict[x] = x
+            files[x] = x
 
-
-class Audio:
-    def add(user,):
-        pass
+        return files
 
 
 @bot.event
@@ -113,13 +157,6 @@ async def on_ready():
 
 @bot.event
 async def on_guild_join(guild):
-    with open('config.json', 'r') as f:
-        prefixes = json.load(f)
-    prefixes[str(guild.id)] = '$'
-
-    with open('config.json', 'r') as f:
-        json.dump(prefixes, f, indent=4)
-
     for g in guild.members:
         DB._fill_bd(g.name, g.id, g.bot, g.nick, guild.id)
 
@@ -128,39 +165,109 @@ async def on_guild_join(guild):
 async def on_member_join(member):
     DB._fill_bd(member.name, member.id, member.bot, member.nick, member.guild.id)
 
+    try:
+        with open('prefix.json', 'r', encoding='utf-8') as fp:  # Open the JSON
+            jsonObject = json.load(fp)  # Load the custom prefixes
+
+    except:
+        jsonObject = {}
+
+    if member.guild.id:  # If the guild exists
+        try:
+            guild_conf = jsonObject[f"{member.guild.id}"]
+            try:
+                bot_role = guild_conf["bot_role"]  # Get bot role
+                guest_role = guild_conf["guest_role"]  # Get guest role
+            except:
+                pass
+        except KeyError:
+            pass
+
+    if bot_role and guest_role:
+        if member.bot == 0:
+            role = discord.utils.get(member.guild.roles, id=guest_role)
+        else:
+            role = discord.utils.get(member.guild.roles, id=bot_role)
+        logging.info(f"Adding to {member} role {role}")
+        await member.add_roles(role)
+
+
+@bot.event
+async def on_member_update(before: discord.Member, after: discord.Member):
+    sql_update_query = f"""UPDATE "{after.guild.id}" set nick = ? where userid = ?"""
+    data_tuple = (after.nick, before.id)
+    DB._work_with_db(sql_update_query, data_tuple)
+
+
+@bot.command(name="upload_audio")
+async def upload_audio(ctx, user=None):
+    user = user or ctx.author
+    if ctx.author.guild_permissions.administrator or user is ctx.author:
+        if ctx.message.attachments:
+            from os import error
+            if not path.isdir(f'tmp/{user.id}'):
+                try:
+                    makedirs(f'tmp/{user.id}')
+                except error as error:
+                    logging.info(f"Failed to create dir", error)
+            if not path.isdir(f'audio/{user.id}'):
+                try:
+                    makedirs(f'audio/{user.id}')
+                except error as error:
+                    logging.info(f"Failed to create dir", error)
+            for at in ctx.message.attachments:
+                import mimetypes
+
+                await at.save(f'tmp/{user.id}/{at.filename}')
+                guess = mimetypes.guess_type(f'tmp/{user.id}/{at.filename}')
+                if guess[0]:
+                    from pymediainfo import MediaInfo
+                    file = f'tmp/{user.id}/{at.filename}'
+                    duration = round(MediaInfo.parse(file).tracks[0].duration / 1000)
+                    if duration > 15:
+                        await ctx.reply(f'Audio duration is {duration}, but max is 15')
+                    else:
+                        a = DB._read_db(ctx.guild.id, user.id)
+                        if a is None:
+                            audiolist = f'{user.id}/{at.filename}'
+                        else:
+                            audiolist = DB._read_db(ctx.guild.id, user.id) + ", " + f'{user.id}/{at.filename}'
+
+                        DB._add_audio(ctx.guild.id, user.id, audiolist)
+                        rename(f'tmp/{user.id}/{at.filename}', f'audio/{user.id}/{at.filename}')
+                else:
+                    await ctx.reply(f'Failed to find MIME type of {at.filename}')
+                    remove(f'tmp/{user.id}/{at.filename}')
+        else:
+            await ctx.reply("Has no Attachment")
+    else:
+        await ctx.reply(f'You`re not admin. You can add audio only for your own account')
+
+
+@bot.command(name="set_prefix")
+@commands.has_permissions(administrator=True)
+async def command_setprefix(ctx, prefix: str):
+    await Commands.setprefix(ctx, prefix)
+
 
 @inter_client.slash_command(
-    name="add-audio",
-    description="Add audio track to user",
+    name="set_prefix",
+    description="Settin up bot prefix",
     options=[
-        Option("audio", "Specify audio name", OptionType.STRING, required=True),
-        Option("user", "Specify any user", OptionType.USER)
+        Option("prefix", "Specify prefix", OptionType.STRING, required=True),
     ]
 )
-async def add_audio(ctx, audio, user=None):
-    user = user or ctx.author
-    a = DB._read_db(ctx.guild.id, user.id)
-    if a == "None":
-        audiolist = audio
-    else:
-        audiolist = DB._read_db(ctx.guild.id, user.id) + ", " + audio
-
-    DB._add_audio(ctx.guild.id, user.id, audiolist)
-    emb = discord.Embed(
-        title=f"Sucseed added {audio} to {user}",
-        color=discord.Color.blue()
-    )
-    emb.set_thumbnail(url=user.avatar_url)
-    await ctx.reply(embed=emb)
+async def slash_setprefix(ctx, prefix: str):
+    await Commands.setprefix(ctx, prefix)
 
 
 @inter_client.slash_command(
     name="info",
     description="Read list of tracks for user",
     options=[
-            Option("user", "Specify any user", OptionType.USER),
-        ]
-    )
+        Option("user", "Specify any user", OptionType.USER),
+    ]
+)
 async def info(ctx, user=None):
     user = user or ctx.author
     audio = DB._read_db(ctx.guild.id, user.id)
@@ -170,12 +277,11 @@ async def info(ctx, user=None):
     else:
         roles = "Not added any role"
 
-    if audio == "None":
+    if audio is None:
         audios = "Not selected audio"
     else:
         audios = "• " + "\n• ".join(sorted(audio.split(", ")))
 
-
     emb = discord.Embed(
         title=f"General information",
         description=f"General information on server about {user}",
@@ -193,91 +299,91 @@ async def info(ctx, user=None):
     await ctx.reply(embed=emb)
 
 
-@bot.command()
-async def select_audio(ctx, user=None):
-    msg = await ctx.send(
-        "Select audio for user!",
-        components=[
-            SelectMenu(
-                custom_id="test",
-                placeholder="List of audio?",
-                max_values=2,
-                options=[
-                    SelectOption("Option 1", "value 1"),
-                    SelectOption("Option 2", "value 2"),
-                    SelectOption("Option 3", "value 3")
-                ]
-            )
-        ]
-    )
-    # Wait for someone to click on it
-    inter = await msg.wait_for_dropdown()
-    # Send what you received
-    labels = [option.label for option in inter.select_menu.selected_options]
-    await inter.reply(f"Options: {', '.join(labels)}")
+@inter_client.slash_command(
+    name="set_trigger_role",
+    description="Setting up role to trigger bot",
+    options=[
+        Option("role", "Specify role", OptionType.ROLE, required=True),
+    ]
+)
+@commands.has_permissions(administrator=True)
+async def set_trigger_role(ctx, role):
+    with open('prefix.json', 'r', encoding='utf-8') as fp:
+        try:
+            jsonObject = json.load(fp)
+        except:
+            jsonObject = {}
+        try:
+            jsonObject[f"{ctx.guild.id}"]["tigger_role"] = role.id
+        except KeyError:
+            trigger_role = jsonObject[f"{ctx.guild.id}"] = {'tigger_role': ''}
+            new = {"tigger_role": role.id}
+            trigger_role.update(new)
+            jsonObject[str(ctx.guild.id)] = trigger_role
+
+    await ctx.send(f"Role to trigger set to : `{role.name}`")
+
+    with open("prefix.json", "w") as fl:
+        json.dump(jsonObject, fl)
+        fl.close()
 
 
-@inter_client.slash_command(description="Select audio from groups to user")
-async def groups(inter):
-    pass
+@inter_client.slash_command(
+    name="set_bot_role",
+    description="Set Default bot role",
+    options=[
+        Option("role", "Specify role", OptionType.ROLE, required=True),
+    ]
+)
+@commands.has_permissions(administrator=True)
+async def set_bot_role(ctx, role):
+    with open('prefix.json', 'r', encoding='utf-8') as fp:
+        try:
+            jsonObject = json.load(fp)
+        except:
+            jsonObject = {}
+        try:
+            jsonObject[f"{ctx.guild.id}"]["bot_role"] = role.id
+        except KeyError:
+            bot_role = jsonObject[f"{ctx.guild.id}"] = {'bot_role': ''}
+            new = {"bot_role": role.id}
+            bot_role.update(new)
+            jsonObject[str(ctx.guild.id)] = bot_role
+
+    await ctx.send(f"Setted up bot role to: `{role.name}`")
+
+    with open("prefix.json", "w") as fl:
+        json.dump(jsonObject, fl)
+        fl.close()
 
 
-@groups.sub_command_group()
-async def select_user_audio(inter, user=None):
-    pass
+@inter_client.slash_command(
+    name="set_guest_role",
+    description="Set Default bot role",
+    options=[
+        Option("role", "Specify role", OptionType.ROLE, required=True),
+    ]
+)
+@commands.has_permissions(administrator=True)
+async def set_guest_role(ctx, role):
+    with open('prefix.json', 'r', encoding='utf-8') as fp:
+        try:
+            jsonObject = json.load(fp)
+        except:
+            jsonObject = {}
+        try:
+            jsonObject[f"{ctx.guild.id}"]["guest_role"] = role.id
+        except KeyError:
+            bot_role = jsonObject[f"{ctx.guild.id}"] = {'guest_role': ''}
+            new = {"guest_role": role.id}
+            bot_role.update(new)
+            jsonObject[str(ctx.guild.id)] = bot_role
 
+    await ctx.send(f"Setted up bot role to: `{role.name}`")
 
-@bot.command(name="upload_audio")
-async def upload_audio(ctx, user=None):
-    user = user or ctx.author
-    if ctx.author.guild_permissions.administrator or user is ctx.author:
-        if ctx.message.attachments:
-            if not os.path.isdir(f'tmp/{user.id}'):
-                try:
-                    os.makedirs(f'tmp/{user.id}')
-                except os.error as error:
-                    logging.info(f"Failed to create dir", error)
-            for a in ctx.message.attachments:
-                import mimetypes
-
-                await a.save(f'tmp/{user.id}/{a.filename}')
-                guess = mimetypes.guess_type(f'tmp/{user.id}/{a.filename}')
-                if guess[0]:
-                    from pymediainfo import MediaInfo
-                    file = f'tmp/{user.id}/{a.filename}'
-                    duration = round(MediaInfo.parse(file).tracks[0].duration / 1000)
-                    if duration > 15:
-                        await ctx.reply(f'Audio duration is {duration}, but max is 10')
-                else:
-                    await ctx.reply(f'Failed to find MIME type of {a.filename}')
-                    os.remove(f'tmp/{user.id}/{a.filename}')
-        else:
-            await ctx.reply("Has no Attachment")
-    else:
-        await ctx.reply(f'You`re not admin. You can add audio only for your own account')
-
-
-@bot.event
-async def on_member_update(before: discord.Member, after: discord.Member):
-    sql_update_query = f"""UPDATE "{after.guild.id}" set nick = ? where userid = ?"""
-    data_tuple = (after.nick, before.id)
-    DB._work_with_db(sql_update_query, data_tuple)
-
-
-@bot.command(pass_context=True)
-@commands.has_permissions(administrator=True) #ensure that only administrators can use this command
-async def changeprefix(ctx, prefix): #command: bl!changeprefix ...
-    with open('config.json', 'r') as f:
-        prefixes = json.load(f)
-
-    prefixes[str(ctx.guild.id)] = prefix
-
-    with open('config.json', 'w') as f: #writes the new prefix into the .json
-        json.dump(prefixes, f, indent=4)
-
-    await ctx.send(f'Prefix changed to: {prefix}') #confirms the prefix it's been changed to
-#next step completely optional: changes bot nickname to also have prefix in the nickname
-    name=f'{prefix}BotBot'
+    with open("prefix.json", "w") as fl:
+        json.dump(jsonObject, fl)
+        fl.close()
 
 
 bot.run('OTQ3OTUzOTAxNzgzNjIxNjYy.GTXbMv.KrztaTO7-ivsPEAVjsyikSQ-GP-ANwULmDraig')