diff --git a/dislocker.py b/dislocker.py index 8065d33..782b287 100644 --- a/dislocker.py +++ b/dislocker.py @@ -18,6 +18,7 @@ class DL(): self.config_dir_path = "./config/" self.export_dir_path = "./export/" self.server_config_path = self.config_dir_path + "server.json" + self.onetime_config_path = self.config_dir_path + "onetime.json" try: if not os.path.isdir(self.config_dir_path): print("config ディレクトリが見つかりません... 作成します。") @@ -48,7 +49,8 @@ class DL(): "search_frequency": 1, "allowable_time": 180, "fstop_time": "21:00:00" - } + }, + "admin_user_id": "TYPE HERE CHANNEL ID (YOU MUST USE INT !!!!)" } } @@ -88,7 +90,7 @@ class DL(): find_pc_list_table = cursor.fetchall() print(find_pc_list_table) if find_pc_list_table[0][0] == False: - cursor.execute("CREATE TABLE pc_list (pc_number INTEGER NOT NULL, using_member_id INTEGER, password_hash VARCHAR(32), PRIMARY KEY (pc_number), FOREIGN KEY (using_member_id) REFERENCES club_member(member_id))") + cursor.execute("CREATE TABLE pc_list (pc_number INTEGER NOT NULL, using_member_id INTEGER, password_hash VARCHAR(32), pc_uuid UUID, pc_token VARCHAR(36), PRIMARY KEY (pc_number), FOREIGN KEY (using_member_id) REFERENCES club_member(member_id))") for i in self.pc_list: print(i) cursor.execute("INSERT INTO pc_list (pc_number) VALUES (%s)", (i,)) @@ -664,8 +666,27 @@ class Bot(discord.Client): pass 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": + if os.path.isfile(dislocker.onetime_config_path): + with open(dislocker.onetime_config_path, "r") as r: + onetime_config = json.load(r) + onetime = str(onetime_config["onetime"]) + await message.channel.send(f"# :dizzy_face: 既にワンタイムパスワードは発行されています。\n# パスワード | {onetime}") + + else: + onetime = str(self.password_generate()) + onetime_config = { + "onetime": str(onetime) + } + 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}") + + """ - msg_split = message.content.split() if msg_split[0] == "/password" or msg_split[0] == "/start": #メッセージの要素が2つ以下の場合は拒否 if len(msg_split) <= 2: diff --git a/dislocker_auth.py b/dislocker_auth.py index 57149b4..b2f10ba 100644 --- a/dislocker_auth.py +++ b/dislocker_auth.py @@ -2,9 +2,13 @@ import psycopg2 import os import json from flask import Flask, request, jsonify, render_template +import uuid +import string +import random config_dir_path = "./config/" server_config_path = config_dir_path + "server.json" +onetime_config_path = config_dir_path + "onetime.json" if not os.path.isfile(server_config_path): if not os.path.isdir(config_dir_path): os.mkdir(config_dir_path) @@ -40,16 +44,38 @@ class Auth(): def __init__(self, host, db, port, user, password): self.db = psycopg2.connect(f"host={host} dbname={db} port={port} user={user} password={password}") - def check(self, pc_number, password): + def token_generate(self, length): + letters = string.ascii_letters + string.digits + password = ''.join(random.choice(letters) for _ in range(length)) + return password + + def check(self, **kwargs): try: cursor = self.db.cursor() - cursor.execute("SELECT * FROM pc_list WHERE pc_number = %s AND password_hash = %s", (pc_number, password)) - pc_info = cursor.fetchall() + pc_number = int(kwargs["pc_number"]) + pc_uuid = uuid.UUID(kwargs["pc_uuid"]) + pc_token = str(kwargs["pc_token"]) + + if "password_hash" in kwargs: + password_hash = str(kwargs["password_hash"]) + cursor.execute("SELECT * FROM pc_list WHERE pc_number = %s AND pc_uuid = %s, AND pc_token = %s", (pc_number, pc_uuid, pc_token)) + pc_info = cursor.fetchall() + else: + cursor.execute("SELECT * FROM pc_list WHERE pc_number = %s AND password_hash = %s AND pc_uuid = %s, AND pc_token = %s", (pc_number, password_hash, pc_uuid, pc_token)) + pc_info = cursor.fetchall() + if pc_info: return {"result": 0, "about": "ok"} else: return {"result": 1, "about": "unregistered_pc"} + except Exception as error: + print("PCの登録状況を調査中にエラーが発生しました。\nエラー内容") + print(str(error.__class__.__name__)) + print(str(error.args)) + print(str(error)) + return {"result": 1, "about": "error"} + finally: cursor.close() @@ -144,17 +170,64 @@ class Auth(): finally: cursor.close() + def register(self, **kwargs): + try: + cursor = self.db.cursor() + pc_number = int(kwargs["pc_number"]) + pc_uuid = uuid.UUID(kwargs["pc_uuid"]) + cursor.execute("SELECT uuid FROM pc_list WHERE pc_number = %s", (pc_number,)) + pc_record = cursor.fetchall() + pc_record_uuid = pc_record[0][0] + if pc_record_uuid == None: + return {"result": 1, "about": "exist"} + else: + pc_token = self.token_generate(36) + cursor.execute("UPDATE pc_list SET pc_uuid = %s, pc_token = %s WHERE pc_number = %s", (pc_uuid, pc_token, pc_number)) + self.db.commit() + os.remove(onetime_config_path) + return {"result": 0, "about": "ok", "output_dict": {"pc_token": pc_token}} + + except Exception as error: + print("停止処理中にエラーが発生しました。\nエラー内容") + print(str(error.__class__.__name__)) + print(str(error.args)) + print(str(error)) + return {"result": 1, "about": "error"} + + finally: + cursor.close() app = Flask(__name__, static_folder="./resource/") auth = Auth(server_config["db"]["host"], server_config["db"]["db_name"], server_config["db"]["port"], server_config["db"]["username"], server_config["db"]["password"]) +@app.route('/register', methods=['POST']) +def register(): + pc_number = int(request.json.get('pc_number')) + pc_uuid = str(request.json.get('pc_uuid')) + onetime_password = int(request.json.get('onetime')) + + if os.path.isfile(onetime_config_path): + with open(onetime_config_path, "r") as r: + onetime_config = json.load(r) + + if onetime_password == onetime_config["onetime"]: + register_result = auth.register(pc_number=pc_number, pc_uuid=pc_uuid) + pc_token = register_result["output_dict"]["pc_token"] + return jsonify({'message': 'ok', 'pc_token': pc_token}), 200 + else: + return jsonify({'message': 'damedesu'}), 401 + else: + return jsonify({'message': 'damedesu'}), 401 + @app.route('/verify', methods=['POST']) def verify(): pc_number = int(request.json.get('pc_number')) - password = request.json.get('password') + password_hash = request.json.get('password') + pc_uuid = request.json.get('pc_uuid') + pc_token = request.json.get('pc_token') print(str(pc_number) + "の認証処理を開始...") - pc_auth = auth.check(pc_number, password) + pc_auth = auth.check(pc_number=pc_number, password_hash=password_hash, pc_uuid=pc_uuid, pc_token=pc_token) if pc_auth["result"] == 0: auth.delete(pc_number) @@ -166,16 +239,23 @@ def verify(): @app.route('/stop', methods=['POST']) def stop(): - pc_number = int(request.json.get('pc_number')) print(str(pc_number) + "の使用停止処理を開始...") - pc_stop = auth.stop(pc_number=pc_number) - if pc_stop["result"] == 0: - print(str(pc_number) + "の使用停止処理は成功しました.") - return jsonify({'message': 'ok'}), 200 - else: - print(str(pc_number) + "の使用停止処理は失敗しました.") - return jsonify({'message': 'error'}), 500 + pc_number = int(request.json.get('pc_number')) + pc_uuid = request.json.get('pc_uuid') + pc_token = request.json.get('pc_token') + pc_auth = auth.check(pc_number=pc_number, pc_uuid=pc_uuid, pc_token=pc_token) + + if pc_auth["result"] == 0: + pc_stop = auth.stop(pc_number=pc_number) + if pc_stop["result"] == 0: + print(str(pc_number) + "の使用停止処理は成功しました.") + return jsonify({'message': 'ok'}), 200 + else: + print(str(pc_number) + "の使用停止処理は失敗しました.") + return jsonify({'message': 'error'}), 500 + else: + return jsonify({'message': 'damedesu'}), 401 if __name__ == '__main__': diff --git a/dislocker_client.py b/dislocker_client.py index 79c5bd7..e6be117 100644 --- a/dislocker_client.py +++ b/dislocker_client.py @@ -11,9 +11,10 @@ import string import random import tkinter import threading -import signal import sys import shutil +import uuid +import time app_name = "Dislocker" dislocker_dir = os.path.dirname(os.path.abspath(sys.argv[0])) @@ -32,13 +33,22 @@ if not os.path.isfile(client_config_path): "auth_host_url": "http://localhost", "pc_number": 1, "master_password_hash": "", - "testing": 0 + "testing": 0, + "pc_uuid": "", + "pc_token": "" } elif os.path.isfile(client_config_path): with open(client_config_path, "r") as r: client_config = json.load(r) +def master_password_gen(): + numbers = string.digits # (1) + password = ''.join(random.choice(numbers) for _ in range(10)) # (2) + password_hash = hashlib.md5(password.encode()).hexdigest() + result = {"password": password, "password_hash": password_hash} + return result + def init(**kwargs): sp_startupinfo = subprocess.STARTUPINFO() sp_startupinfo.dwFlags = subprocess.STARTF_USESHOWWINDOW @@ -52,17 +62,46 @@ def init(**kwargs): return 1 if client_config["initial"] == 1: - master_password = master_password_gen() - msgbox = tkinter.messagebox.showinfo(title=f"{app_name} | 初回起動を検出", message=f"初回起動のようです。\nマスターパスワードを記録しておいてください。\nこれ以降二度と表示されることはないでしょう。\n\n{master_password["password"]}\n\nまた、認証先サーバーの接続先を指定してください。ロックを解除できなくなります。") - client_config["master_password_hash"] = master_password["password_hash"] - client_config["initial"] = 0 + pc_uuid = uuid.uuid4() + client_config["pc_uuid"] = str(pc_uuid) + if "pc_number" in kwargs: client_config["pc_number"] = int(kwargs["pc_number"]) else: - client_config["pc_number"] = 1 - with open(client_config_path, "w") as w: - json.dump(client_config, w, indent=4) - return 2 + tkinter.messagebox.showerror(title=f"{app_name} | 登録時にエラー", message=f"登録時にエラーが発生しました。\nPC番号が指定されていません。1個目の引数にPC番号、2個目の引数にワンタイムパスワードを指定して、もう一度お試しください。") + return 1 + + if "onetime" in kwargs: + onetime = str(kwargs["onetime"]) + else: + tkinter.messagebox.showerror(title=f"{app_name} | 登録時にエラー", message=f"登録時にエラーが発生しました。\nワンタイムパスワードが指定されていません。1個目の引数にPC番号、2個目の引数にワンタイムパスワードを指定して、もう一度お試しください。") + return 1 + + register_url = client_config["auth_host_url"] + "/register" + register_json = { + "pc_number": int(client_config["pc_number"]), + "pc_uuid": str(pc_uuid), + "onetime": onetime + } + + responce = requests.post(register_url, json=register_json) + + if responce.status_code == 200: + print("PCの情報が登録されました。") + pc_token = str(responce.json["pc_token"]) + client_config["pc_token"] = pc_token + + master_password = master_password_gen() + msgbox = tkinter.messagebox.showinfo(title=f"{app_name} | 初回起動を検出", message=f"初回起動のようです。\nマスターパスワードを記録しておいてください。\nこれ以降二度と表示されることはないでしょう。\n\n{master_password["password"]}\n\nまた、認証先サーバーの接続先を指定してください。ロックを解除できなくなります。") + client_config["master_password_hash"] = master_password["password_hash"] + client_config["initial"] = 0 + + with open(client_config_path, "w") as w: + json.dump(client_config, w, indent=4) + return 2 + else: + msgbox = tkinter.messagebox.showerror(title=f"{app_name} | 登録時にエラー", message=f"登録時にエラーが発生しました。\nワンタイムパスワードが間違っている可能性があります。") + return 1 else: return 0 @@ -89,37 +128,6 @@ class App(customtkinter.CTk): self.unlock_taskmgr() self.toast() self.destroy() - - def delete_appdata(self, **kwargs): - process_name = kwargs["process_name"] - dir_path = kwargs["dir_path"] - - if not os.path.exists(dir_path): - print(f"エラー: 指定されたディレクトリ {dir_path} が存在しません。") - return 1 - - try: - # プロセスの終了 - subprocess.run(['taskkill', '/f', '/t', '/im', process_name]) - print(f"{process_name} を終了しました。") - - # ディレクトリの削除 - shutil.rmtree(dir_path) - print(f"{dir_path} を削除しました。") - - return 0 - except subprocess.CalledProcessError as e: - print(f"プロセス終了エラー: {e}") - return 1 - except PermissionError as e: - print(f"権限エラー: {e}") - return 1 - except Exception as error: - print("エラーが発生しました。\nエラー内容:") - print(f"エラータイプ: {error.__class__.__name__}") - print(f"エラー引数: {error.args}") - print(f"エラーメッセージ: {str(error)}") - return 1 def block_key(self): block_keys = ['ctrl', 'alt', 'windows', 'shift', 'delete'] @@ -263,6 +271,7 @@ class Lock(customtkinter.CTkToplevel): def auth(self): self.button_disable() password = str(self.password_entry.get()) + if len(password) == 10: print("マスターパスワードで認証を試行します。") master_password_hash = self.hash_genarate(str(self.password_entry.get())) @@ -282,6 +291,8 @@ class Lock(customtkinter.CTkToplevel): auth_url = client_config["auth_host_url"] + "/verify" auth_json = { "pc_number": int(client_config["pc_number"]), + "pc_uuid": str(client_config["pc_uuid"]), + "pc_token": str(client_config["pc_token"]), "password": self.hash_genarate(str(self.password_entry.get())) } try: @@ -349,17 +360,73 @@ class Help(customtkinter.CTkToplevel): def handler_close(self): self.destroy() - -class Monitor(): + + +class Stop(): def __init__(self) -> None: pass - def start(self): - monitor_thread = threading.Thread(target=self.run) - monitor_thread.start() + def run(self): + stop_thread = threading.Thread(target=self.stop) + stop_thread.run() + run_notify = Notification( + app_id='Dislocker', + title='終了処理を実行中', + msg='終了処理を実行しています。\nPCがシャットダウンするまで、そのままでお待ち下さい。', + icon=resource_path + r'\success.png' + ) + run_notify.set_audio(audio.Default, loop=False) + run_notify.show() - def signal_handler(self): + def delete_appdata(self, **kwargs): + process_name = kwargs["process_name"] + dir_path = kwargs["dir_path"] + + if not os.path.exists(dir_path): + print(f"エラー: 指定されたディレクトリ {dir_path} が存在しません。") + return 1 + + try: + # プロセスの終了 + subprocess.run(['taskkill', '/f', '/t', '/im', process_name]) + print(f"{process_name} を終了しました。") + + time.sleep(0.1) + + # ディレクトリの削除 + i = 1 + ic = 0 + while i == 1: + shutil.rmtree(dir_path) + if os.path.isdir(dir_path): + ic += 1 + if ic == 10: + i = 0 + else: + i = 0 + print(f"{dir_path} を削除しました。") + + return 0 + + except subprocess.CalledProcessError as e: + print(f"プロセス終了エラー: {e}") + return 1 + except PermissionError as e: + print(f"権限エラー: {e}") + return 1 + except Exception as error: + print("エラーが発生しました。\nエラー内容:") + print(f"エラータイプ: {error.__class__.__name__}") + print(f"エラー引数: {error.args}") + print(f"エラーメッセージ: {str(error)}") + return 1 + + def shutdown(self): + shutdown_command = subprocess.run(['shutdown', '/s', '/t', '1']) + + def stop(self): print("停止処理を実行。") + appdata_local = os.path.expandvars("%LOCALAPPDATA%") appdata_roaming = os.path.expandvars("%APPDATA%") epic_del = app.delete_appdata(process_name="EpicGamesLauncher.exe", dir_path=f"{appdata_local}\\EpicGamesLauncher\\Saved") @@ -370,12 +437,17 @@ class Monitor(): riot_del = app.delete_appdata(process_name="RiotClientServices.exe", dir_path=f"{appdata_local}\\Riot Games\\Riot Client") stop_url = client_config["auth_host_url"] + "/stop" stop_json = { - "pc_number": int(client_config["pc_number"]) + "pc_number": int(client_config["pc_number"]), + "pc_uuid": str(client_config["pc_uuid"]), + "pc_token": str(client_config["pc_token"]) } try: responce = requests.post(stop_url, json=stop_json) if responce.status_code == 200: print("停止処理は成功しました。") + elif responce.status_code == 401: + print("認証に失敗しました。") + tkinter.messagebox.showwarning(title=f"{app_name} | エラー", message=f"認証に失敗しました。\nDiscordサーバーの指示に従って、停止処理を自身で行ってください。") else: print("内部エラーにより停止処理に失敗しました。") result_msgbox = tkinter.messagebox.showwarning(title=f"{app_name} | エラー", message=f"内部エラーにより停止処理に失敗しました。\nDiscordサーバーの指示に従って、停止処理を自身で行ってください。") @@ -385,21 +457,6 @@ class Monitor(): finally: self.shutdown() - def shutdown(self): - shutdown_command = subprocess.run(['shutdown', '/s', '/t', '1']) - - - def run(self): - signal.signal(signal.SIGTERM, self.signal_handler) - - -def master_password_gen(): - numbers = string.digits # (1) - password = ''.join(random.choice(numbers) for _ in range(10)) # (2) - password_hash = hashlib.md5(password.encode()).hexdigest() - result = {"password": password, "password_hash": password_hash} - return result - if __name__ == '__main__': args = sys.argv @@ -411,11 +468,11 @@ if __name__ == '__main__': elif init_result == 2: pass else: - app = App() - monitor = Monitor() - monitor.signal_handler() + stop = Stop() + stop.run() + elif args[1] == "setup": - init_result = init(pc_number=args[2]) + init_result = init(pc_number=args[2], onetime=args[3]) if init_result == 1: warning_msgbox = tkinter.messagebox.showwarning(title=f"{app_name} | 多重起動エラー", message=f"すでに {app_name} は実行されています。\n正常に起動しない場合は、既に起動しているプロセスを終了してから、もう一度起動してみてください。") elif init_result == 2: