diff --git a/dislocker_slash.py b/dislocker_slash.py index 39f1b47..4f7893b 100644 --- a/dislocker_slash.py +++ b/dislocker_slash.py @@ -645,62 +645,107 @@ class DL(): if cursor: cursor.close() - def pc_onetime_gen(self): + def pc_onetime_gen(self, **kwargs): + if kwargs.get("max_count") == None: + max_count = 1 + elif kwargs.get("max_count").isdecimal(): + max_count = int(kwargs.get("max_count")) + else: + max_count = 1 + try: if os.path.isfile(dislocker.onetime_config_path): with open(dislocker.onetime_config_path, "r") as r: onetime_config = json.load(r) - onetime = onetime_config["onetime"]["pc_register"] - if onetime == None: - onetime = str(self.password_generate(8)) - onetime_config["onetime"]["pc_register"] = onetime + onetime_password = onetime_config["onetime"]["pc_register"]["password"] + if onetime_password == None: + onetime_password = str(self.password_generate(8)) + onetime_config["onetime"]["pc_register"]["password"] = onetime_password + onetime_config["onetime"]["pc_register"]["max_count"] = max_count + onetime_config["onetime"]["pc_register"]["current_count"] = 0 + current_count = onetime_config["onetime"]["pc_register"]["current_count"] with open(dislocker.onetime_config_path, "w") as w: json.dump(onetime_config, w, indent=4) - return {"result": 0, "about": "ok", "onetime": onetime} + return {"result": 0, "about": "ok", "output_dict": {"onetime_password": onetime_password, "current_count": current_count, "max_count": max_count}} else: - return {"result": 0, "already_exists": "ok", "onetime": onetime} + current_count = onetime_config["onetime"]["pc_register"]["current_count"] + max_count = onetime_config["onetime"]["pc_register"]["max_count"] + return {"result": 1, "about": "already_exists", "output_dict": {"onetime_password": onetime_password, "current_count": current_count, "max_count": max_count}} else: - onetime = str(self.password_generate(8)) + onetime_password = str(self.password_generate(8)) onetime_config = { "onetime": { - "pc_register": onetime, - "device_register": None + "pc_register": { + "password": onetime_password, + "current_count": 0, + "max_count": int(max_count) + }, + "device_register": { + "password": None, + "current_count": None, + "max_count": None + } } } + current_count = onetime_config["onetime"]["pc_register"]["current_count"] with open(dislocker.onetime_config_path, "w") as w: json.dump(onetime_config, w, indent=4) - return {"result": 0, "about": "ok", "onetime": onetime} + return {"result": 0, "about": "ok", "output_dict": {"onetime_password": onetime_password, "current_count": current_count, "max_count": max_count}} + except Exception as error: - self.log(title=f"[ERROR] PC登録用のワンタイムパスワードを発行中にエラーが発生しました。 {str(error.__class__.__name__)}", message=str(error.args), flag=1) + self.log(title=f"[ERROR] PC登録用のワンタイムパスワード発行中にエラーが発生しました。 {str(error.__class__.__name__)}", message=str(error.args), flag=1) return {"result": 1, "about": "error"} + - def device_onetime_gen(self): + def device_onetime_gen(self, **kwargs): + if kwargs.get("max_count") == None: + max_count = 1 + elif kwargs.get("max_count").isdecimal(): + max_count = int(kwargs.get("max_count")) + else: + max_count = 1 + try: if os.path.isfile(dislocker.onetime_config_path): with open(dislocker.onetime_config_path, "r") as r: onetime_config = json.load(r) - onetime = onetime_config["onetime"]["device_register"] - if onetime == None: - onetime = str(self.password_generate(8)) - onetime_config["onetime"]["device_register"] = onetime + onetime_password = onetime_config["onetime"]["device_register"]["password"] + if onetime_password == None: + onetime_password = str(self.password_generate(8)) + onetime_config["onetime"]["device_register"]["password"] = onetime_password + onetime_config["onetime"]["device_register"]["max_count"] = max_count + onetime_config["onetime"]["device_register"]["current_count"] = 0 + current_count = onetime_config["onetime"]["device_register"]["current_count"] with open(dislocker.onetime_config_path, "w") as w: json.dump(onetime_config, w, indent=4) - return {"result": 0, "about": "ok", "onetime": onetime} + return {"result": 0, "about": "ok", "output_dict": {"onetime_password": onetime_password, "current_count": current_count, "max_count": max_count}} else: - return {"result": 0, "already_exists": "ok", "onetime": onetime} + current_count = onetime_config["onetime"]["device_register"]["current_count"] + max_count = onetime_config["onetime"]["device_register"]["max_count"] + return {"result": 1, "about": "already_exists", "output_dict": {"onetime_password": onetime_password, "current_count": current_count, "max_count": max_count}} else: - onetime = str(self.password_generate(8)) + onetime_password = str(self.password_generate(8)) onetime_config = { "onetime": { - "pc_register": None, - "device_register": onetime + "pc_register": { + "password": None, + "current_count": None, + "max_count": None + }, + "device_register": { + "password": onetime_password, + "current_count": 0, + "max_count": int(max_count) + } } } + current_count = onetime_config["onetime"]["device_register"]["current_count"] with open(dislocker.onetime_config_path, "w") as w: json.dump(onetime_config, w, indent=4) - return {"result": 0, "about": "ok", "onetime": onetime} + return {"result": 0, "about": "ok", "output_dict": {"onetime_password": onetime_password, "current_count": current_count, "max_count": max_count}} + except Exception as error: - self.log(title=f"[ERROR] デバイス登録用のワンタイムパスワードを発行中にエラーが発生しました。 {str(error.__class__.__name__)}", message=str(error.args), flag=1) + self.log(title=f"[ERROR] PC登録用のワンタイムパスワード発行中にエラーが発生しました。 {str(error.__class__.__name__)}", message=str(error.args), flag=1) return {"result": 1, "about": "error"} @@ -881,87 +926,61 @@ async def on_message(self, message): elif isinstance(message.channel, discord.DMChannel): if message.author.id == dislocker.server_config["bot"]["admin_user_id"]: msg_split = message.content.split() + if msg_split[0] == "/pcreg": max_count = 1 if len(msg_split) == 2: if msg_split[1].isdecimal(): max_count = int(msg_split[1]) + + pc_onetime_password_gen = dislocker.pc_onetime_gen(max_count=max_count) + + if pc_onetime_password_gen["result"] == 0: + pc_onetime_password = str(pc_onetime_password_gen["output_dict"]["onetime_password"]) + pc_onetime_password_max_count = str(pc_onetime_password_gen["output_dict"]["max_count"]) + pc_onetime_password_current_count = str(pc_onetime_password_gen["output_dict"]["current_count"]) + pc_onetime_password_remaining_times = str(int(pc_onetime_password_max_count) - int(pc_onetime_password_current_count)) - if os.path.isfile(dislocker.onetime_config_path): - with open(dislocker.onetime_config_path, "r") as r: - onetime_config = json.load(r) - onetime = onetime_config["onetime"]["pc_register"]["password"] - if onetime == None: - onetime = str(self.password_generate(8)) - onetime_config["onetime"]["pc_register"]["password"] = onetime - onetime_config["onetime"]["pc_register"]["max_count"] = max_count - with open(dislocker.onetime_config_path, "w") as w: - json.dump(onetime_config, w, indent=4) - await message.channel.send(f"# :dizzy_face: PC登録時のワンタイムパスワードを発行します。\n# パスワード | {onetime}\n# 最大使用回数 | {str(max_count)}") + await message.channel.send(f"# :dizzy_face: PC登録時のワンタイムパスワードを発行します。\n# パスワード | {pc_onetime_password}\n# 最大使用回数 | {pc_onetime_password_max_count}\n# 残り使用回数 | {pc_onetime_password_remaining_times}") + elif pc_onetime_password_gen["result"] == 1: + if pc_onetime_password_gen["about"] == "already_exists": + pc_onetime_password = str(pc_onetime_password_gen["output_dict"]["onetime_password"]) + pc_onetime_password_max_count = str(pc_onetime_password_gen["output_dict"]["max_count"]) + pc_onetime_password_current_count = str(pc_onetime_password_gen["output_dict"]["current_count"]) + pc_onetime_password_remaining_times = str(int(pc_onetime_password_max_count) - int(pc_onetime_password_current_count)) + await message.channel.send(f"# :dizzy_face: 既にワンタイムパスワードは発行されています。\n# パスワード | {pc_onetime_password}\n# 残り使用回数 | {pc_onetime_password_remaining_times}") else: - await message.channel.send(f"# :dizzy_face: 既にワンタイムパスワードは発行されています。\n# パスワード | {onetime}\n# 残り使用回数 | {str(max_count - onetime_config['onetime']['pc_register']['current_count'])}") - else: - onetime = str(self.password_generate(8)) - onetime_config = { - "onetime": { - "pc_register": { - "password": onetime, - "current_count": 0, - "max_count": int(max_count) - }, - "device_register": { - "password": None, - "current_count": None, - "max_count": None - } - } - } - with open(dislocker.onetime_config_path, "w") as w: - json.dump(onetime_config, w, indent=4) - await message.channel.send(f"# :dizzy_face: PC登録時のワンタイムパスワードを発行します。\n# パスワード | {onetime}\n# 最大使用回数 | {str(max_count)}") + await message.channel.send("# :skull_crossbones: ワンタイムパスワードの発行に失敗しました。") + elif msg_split[0] == "/devreg": max_count = 1 if len(msg_split) == 2: if msg_split[1].isdecimal(): max_count = int(msg_split[1]) - if os.path.isfile(dislocker.onetime_config_path): - with open(dislocker.onetime_config_path, "r") as r: - onetime_config = json.load(r) - onetime = onetime_config["onetime"]["device_register"]["password"] - if onetime == None: - onetime = str(self.password_generate(8)) - onetime_config["onetime"]["device_register"]["password"] = onetime - onetime_config["onetime"]["device_register"]["max_count"] = max_count - with open(dislocker.onetime_config_path, "w") as w: - json.dump(onetime_config, w, indent=4) - await message.channel.send(f"# :dizzy_face: デバイス登録時のワンタイムパスワードを発行します。\n# パスワード | {onetime}\n# 最大使用回数 | {str(max_count)}") + device_onetime_password_gen = dislocker.device_onetime_gen(max_count=max_count) + + if device_onetime_password_gen["result"] == 0: + device_onetime_password = str(device_onetime_password_gen["output_dict"]["onetime_password"]) + device_onetime_password_max_count = str(device_onetime_password_gen["output_dict"]["max_count"]) + device_onetime_password_current_count = str(device_onetime_password_gen["output_dict"]["current_count"]) + device_onetime_password_remaining_times = str(int(device_onetime_password_max_count) - int(device_onetime_password_current_count)) + + await message.channel.send(f"# :dizzy_face: デバイス登録時のワンタイムパスワードを発行します。\n# パスワード | {device_onetime_password}\n# 最大使用回数 | {device_onetime_password_max_count}\n# 残り使用回数 | {device_onetime_password_remaining_times}") + elif device_onetime_password_gen["result"] == 1: + if device_onetime_password_gen["about"] == "already_exists": + device_onetime_password = str(device_onetime_password_gen["output_dict"]["onetime_password"]) + device_onetime_password_max_count = str(device_onetime_password_gen["output_dict"]["max_count"]) + device_onetime_password_current_count = str(device_onetime_password_gen["output_dict"]["current_count"]) + device_onetime_password_remaining_times = str(int(device_onetime_password_max_count) - int(device_onetime_password_current_count)) + await message.channel.send(f"# :dizzy_face: 既にワンタイムパスワードは発行されています。\n# パスワード | {device_onetime_password}\n# 残り使用回数 | {device_onetime_password_remaining_times}") else: - await message.channel.send(f"# :dizzy_face: 既にワンタイムパスワードは発行されています。\n# パスワード | {onetime}\n# 残り使用回数 | {str(max_count - onetime_config['onetime']['device_register']['current_count'])}") - else: - onetime = str(self.password_generate(8)) - onetime_config = { - "onetime": { - "pc_register": { - "password": None, - "current_count": None, - "max_count": None - }, - "device_register": { - "password": onetime, - "current_count": 0, - "max_count": int(max_count) - } - } - } - with open(dislocker.onetime_config_path, "w") as w: - json.dump(onetime_config, w, indent=4) - await message.channel.send(f"# :dizzy_face: デバイス登録時のワンタイムパスワードを発行します。\n# パスワード | {onetime}\n# 最大使用回数 | {str(max_count)}") + await message.channel.send("# :skull_crossbones: ワンタイムパスワードの発行に失敗しました。") else: await message.channel.send("# :warning: DMでの応答は、現在無効化されています。") else: - await message.channel.send("# :warning: DMでの応答は、現在無効化されています。") + pass @client.event async def on_interaction(interaction: discord.Interaction): @@ -1175,27 +1194,53 @@ async def userreg(interaction: discord.Interaction, discord_user_id: str, discor elif user_register["about"] == "error": await interaction.response.send_message(":x: 内部エラーが発生しました。\nサーバーでエラーが発生しています。管理者に問い合わせてください。", ephemeral=True) -@tree.command(name="pcreg", description="PCを登録するためのワンタイムパスワードを発行します。") +@tree.command(name="pcreg", description="PCをDislockerに登録するためのワンタイムパスワードを発行します。") @discord.app_commands.default_permissions(administrator=True) -async def pcreg(interaction: discord.Interaction): - onetime = dislocker.pc_onetime_gen() - if onetime["result"] == 0: - await interaction.response.send_message(f":white_check_mark: PC登録用のワンタイムパスワードを発行しました。\n>>> # ワンタイムパスワード | {onetime['onetime']}\n## 有効期限 | 1回の登録でのみ使用可能です。", ephemeral=True) - dislocker.log(title=f"[INFO] PC登録用のワンタイムパスワードを発行しました。", message=f"ワンタイムパスワード | {onetime['onetime']}", flag=0) - elif onetime["result"] == 1: - if onetime["about"] == "error": - await interaction.response.send_message(":x: 内部エラーが発生しました。\nサーバーでエラーが発生しています。管理者に問い合わせてください。", ephemeral=True) +async def pcreg(interaction: discord.Interaction, how_much: int = 1): + max_count = how_much + + pc_onetime_password_gen = dislocker.pc_onetime_gen(max_count=max_count) + + if pc_onetime_password_gen["result"] == 0: + pc_onetime_password = str(pc_onetime_password_gen["output_dict"]["onetime_password"]) + pc_onetime_password_max_count = str(pc_onetime_password_gen["output_dict"]["max_count"]) + pc_onetime_password_current_count = str(pc_onetime_password_gen["output_dict"]["current_count"]) + pc_onetime_password_remaining_times = str(int(pc_onetime_password_max_count) - int(pc_onetime_password_current_count)) -@tree.command(name="devicereg", description="デバイスを登録するためのワンタイムパスワードを発行します。") + await interaction.response.send_message(f"# :dizzy_face: PC登録時のワンタイムパスワードを発行します。\n# パスワード | {pc_onetime_password}\n# 最大使用回数 | {pc_onetime_password_max_count}\n# 残り使用回数 | {pc_onetime_password_remaining_times}", ephemeral=True) + elif pc_onetime_password_gen["result"] == 1: + if pc_onetime_password_gen["about"] == "already_exists": + pc_onetime_password = str(pc_onetime_password_gen["output_dict"]["onetime_password"]) + pc_onetime_password_max_count = str(pc_onetime_password_gen["output_dict"]["max_count"]) + pc_onetime_password_current_count = str(pc_onetime_password_gen["output_dict"]["current_count"]) + pc_onetime_password_remaining_times = str(int(pc_onetime_password_max_count) - int(pc_onetime_password_current_count)) + await interaction.response.send_message(f"# :dizzy_face: 既にワンタイムパスワードは発行されています。\n# パスワード | {pc_onetime_password}\n# 残り使用回数 | {pc_onetime_password_remaining_times}", ephemeral=True) + else: + await interaction.response.send_message("# :skull_crossbones: ワンタイムパスワードの発行に失敗しました。", ephemeral=True) + +@tree.command(name="devicereg", description="デバイスをDislockerに登録するためのワンタイムパスワードを発行します。") @discord.app_commands.default_permissions(administrator=True) -async def devicereg(interaction: discord.Interaction): - onetime = dislocker.device_onetime_gen() - if onetime["result"] == 0: - await interaction.response.send_message(f":white_check_mark: デバイス登録用のワンタイムパスワードを発行しました。\n>>> # ワンタイムパスワード | {onetime['onetime']}\n## 有効期限 | 1回の登録でのみ使用可能です。", ephemeral=True) - dislocker.log(title=f"[INFO] デバイス登録用のワンタイムパスワードを発行しました。", message=f"ワンタイムパスワード | {onetime['onetime']}", flag=0) - elif onetime["result"] == 1: - if onetime["about"] == "error": - await interaction.response.send_message(":x: 内部エラーが発生しました。\nサーバーでエラーが発生しています。管理者に問い合わせてください。", ephemeral=True) +async def devicereg(interaction: discord.Interaction, how_much: int): + max_count = how_much + + device_onetime_password_gen = dislocker.device_onetime_gen(max_count=max_count) + + if device_onetime_password_gen["result"] == 0: + device_onetime_password = str(device_onetime_password_gen["output_dict"]["onetime_password"]) + device_onetime_password_max_count = str(device_onetime_password_gen["output_dict"]["max_count"]) + device_onetime_password_current_count = str(device_onetime_password_gen["output_dict"]["current_count"]) + device_onetime_password_remaining_times = str(int(device_onetime_password_max_count) - int(device_onetime_password_current_count)) + + await interaction.response.send_message(f"# :dizzy_face: デバイス登録時のワンタイムパスワードを発行します。\n# パスワード | {device_onetime_password}\n# 最大使用回数 | {device_onetime_password_max_count}\n# 残り使用回数 | {device_onetime_password_remaining_times}", ephemeral=True) + elif device_onetime_password_gen["result"] == 1: + if device_onetime_password_gen["about"] == "already_exists": + device_onetime_password = str(device_onetime_password_gen["output_dict"]["onetime_password"]) + device_onetime_password_max_count = str(device_onetime_password_gen["output_dict"]["max_count"]) + device_onetime_password_current_count = str(device_onetime_password_gen["output_dict"]["current_count"]) + device_onetime_password_remaining_times = str(int(device_onetime_password_max_count) - int(device_onetime_password_current_count)) + await interaction.response.send_message(f"# :dizzy_face: 既にワンタイムパスワードは発行されています。\n# パスワード | {device_onetime_password}\n# 残り使用回数 | {device_onetime_password_remaining_times}", ephemeral=True) + else: + await interaction.response.send_message("# :skull_crossbones: ワンタイムパスワードの発行に失敗しました。", ephemeral=True) @tree.command(name="fstop", description="PCの使用登録を強制的に終了します。") @discord.app_commands.default_permissions(administrator=True)