From 8ac3f498ed62103f6017422143d26f251306a378 Mon Sep 17 00:00:00 2001 From: Karl Date: Tue, 15 Jul 2025 15:45:17 +0100 Subject: [PATCH] docstrings --- debug_account.py | 12 ++-- ktvmanager/account_checker.py | 36 ++++++++--- ktvmanager/lib/auth.py | 55 ++++++++++++++--- ktvmanager/lib/checker.py | 85 +++++++++++++++++++------- ktvmanager/lib/database.py | 109 ++++++++++++++++++++++++++++------ ktvmanager/lib/get_urls.py | 38 ++++++++---- routes/api.py | 108 +++++++++++++++++++++++++++++---- 7 files changed, 362 insertions(+), 81 deletions(-) diff --git a/debug_account.py b/debug_account.py index 9088e9c..ec1be19 100644 --- a/debug_account.py +++ b/debug_account.py @@ -11,7 +11,9 @@ sys.path.append(project_root) from ktvmanager.lib.checker import single_account_check from ktvmanager.lib.get_urls import get_latest_urls_from_dns -def main(): + +def main() -> None: + """Debugs a single KTV account by checking its validity against available stream URLs.""" parser = argparse.ArgumentParser(description="Debug a single KTV account.") parser.add_argument("username", help="The username to check.") parser.add_argument("password", help="The password to check.") @@ -24,15 +26,12 @@ def main(): user=os.getenv("DBUSER"), password=os.getenv("DBPASS"), database=os.getenv("DATABASE"), - port=os.getenv("DBPORT") + port=os.getenv("DBPORT"), ) stream_urls = get_latest_urls_from_dns() - account_data = { - "username": args.username, - "password": args.password - } + account_data = {"username": args.username, "password": args.password} result = single_account_check(account_data, stream_urls) @@ -43,5 +42,6 @@ def main(): db_connection.close() + if __name__ == "__main__": main() \ No newline at end of file diff --git a/ktvmanager/account_checker.py b/ktvmanager/account_checker.py index 1fcec19..532d389 100644 --- a/ktvmanager/account_checker.py +++ b/ktvmanager/account_checker.py @@ -9,8 +9,19 @@ sys.path.append(project_root) from ktvmanager.lib.encryption import decrypt_password from ktvmanager.lib.checker import single_account_check from ktvmanager.lib.get_urls import get_latest_urls_from_dns +from typing import List, Dict, Any +from mysql.connector.connection import MySQLConnection -def get_all_accounts(db_connection): + +def get_all_accounts(db_connection: MySQLConnection) -> List[Dict[str, Any]]: + """Retrieves all user accounts from the database. + + Args: + db_connection: An active MySQL database connection. + + Returns: + A list of dictionaries, where each dictionary represents a user account. + """ cursor = db_connection.cursor(dictionary=True) query = "SELECT * FROM userAccounts where userId = 1" cursor.execute(query) @@ -18,7 +29,11 @@ def get_all_accounts(db_connection): cursor.close() return accounts -def main(): + +def main() -> None: + """ + Checks the validity of all accounts in the database against available stream URLs. + """ load_dotenv() db_connection = mysql.connector.connect( @@ -26,7 +41,7 @@ def main(): user=os.getenv("DBUSER"), password=os.getenv("DBPASS"), database=os.getenv("DATABASE"), - port=os.getenv("DBPORT") + port=os.getenv("DBPORT"), ) accounts = get_all_accounts(db_connection) @@ -34,24 +49,29 @@ def main(): for account in accounts: try: - decrypted_password = decrypt_password(account['password']) + decrypted_password = decrypt_password(account["password"]) except Exception as e: print(f"Could not decrypt password for {account['username']}: {e}") continue account_data = { - "username": account['username'], - "password": decrypted_password + "username": account["username"], + "password": decrypted_password, } result = single_account_check(account_data, stream_urls) if result: - print(f"Account {account['username']} on stream {account['stream']} is VALID.") + print( + f"Account {account['username']} on stream {account['stream']} is VALID." + ) else: - print(f"Account {account['username']} on stream {account['stream']} is INVALID.") + print( + f"Account {account['username']} on stream {account['stream']} is INVALID." + ) db_connection.close() + if __name__ == "__main__": main() \ No newline at end of file diff --git a/ktvmanager/lib/auth.py b/ktvmanager/lib/auth.py index f62133c..355e488 100644 --- a/ktvmanager/lib/auth.py +++ b/ktvmanager/lib/auth.py @@ -1,19 +1,60 @@ from functools import wraps -from flask import request, jsonify, Blueprint, make_response +from flask import request, jsonify, Blueprint, Response +from typing import Callable, Any, Tuple + auth_blueprint = Blueprint("auth", __name__) -def check_auth(username, password): + +def check_auth(username: str, password: str) -> bool: + """ + This function checks if a username and password are valid. + Currently, it always returns True. + + Args: + username: The username to check. + password: The password to check. + + Returns: + True if the credentials are valid, False otherwise. + """ return True -def requires_basic_auth(f): + +def requires_basic_auth(f: Callable) -> Callable: + """ + A decorator to protect routes with basic authentication. + + Args: + f: The function to decorate. + + Returns: + The decorated function. + """ + @wraps(f) - def decorated(*args, **kwargs): + def decorated(*args: Any, **kwargs: Any) -> Tuple[Response, int, Dict[str, str]] | Response: auth = request.authorization if not auth or not check_auth(auth.username, auth.password): - return jsonify({"message": "Could not verify"}), 401, {'WWW-Authenticate': 'Basic realm="Login Required"'} - return f(auth.username,auth.password, *args, **kwargs) + return ( + jsonify({"message": "Could not verify"}), + 401, + {"WWW-Authenticate": 'Basic realm="Login Required"'}, + ) + return f(auth.username, auth.password, *args, **kwargs) + return decorated -def check_login(username, password): + +def check_login(username: str, password: str) -> Response: + """ + Checks a user's login credentials. + + Args: + username: The username to check. + password: The password to check. + + Returns: + A Flask JSON response indicating success. + """ return jsonify({"auth": "Success"}) \ No newline at end of file diff --git a/ktvmanager/lib/checker.py b/ktvmanager/lib/checker.py index dca6fcd..1ec8209 100644 --- a/ktvmanager/lib/checker.py +++ b/ktvmanager/lib/checker.py @@ -1,12 +1,34 @@ import requests from concurrent.futures import ThreadPoolExecutor, as_completed -from flask import request, jsonify +from typing import List, Optional, Dict, Any, Tuple +from flask import request, jsonify, Response from ktvmanager.lib.get_urls import get_latest_urls_from_dns -def build_url(stream_url, username, password): + +def build_url(stream_url: str, username: str, password: str) -> str: + """Builds the player API URL for a given stream URL and credentials. + + Args: + stream_url: The base URL of the streaming server. + username: The username for the account. + password: The password for the account. + + Returns: + The fully constructed player API URL. + """ return f"{stream_url}/player_api.php?username={username}&password={password}" -def check_url(url): + +def check_url(url: str) -> Optional[Dict[str, Any]]: + """Checks if a given URL is a valid and authenticated streaming service endpoint. + + Args: + url: The URL to check. + + Returns: + A dictionary containing the JSON response from the server if the account is + valid, otherwise None. + """ try: response = requests.get(url, timeout=5) response.raise_for_status() @@ -18,30 +40,49 @@ def check_url(url): return None return None -def single_account_check(account_data, stream_urls): + +def single_account_check( + account_data: Dict[str, str], stream_urls: List[str] +) -> Optional[Dict[str, Any]]: + """Checks a single account against a list of stream URLs concurrently. + + Args: + account_data: A dictionary containing the 'username' and 'password'. + stream_urls: A list of stream URLs to check against. + + Returns: + A dictionary containing the valid URL and the server's response data + if a valid URL is found, otherwise None. + """ if not stream_urls: return None - executor = ThreadPoolExecutor(max_workers=min(10, len(stream_urls))) - future_to_url = { - executor.submit( - check_url, - build_url(stream_url, account_data['username'], account_data['password']) - ): stream_url - for stream_url in stream_urls - } + with ThreadPoolExecutor(max_workers=min(10, len(stream_urls))) as executor: + future_to_url = { + executor.submit( + check_url, + build_url( + stream_url, account_data["username"], account_data["password"] + ), + ): stream_url + for stream_url in stream_urls + } - final_result = None - for future in as_completed(future_to_url): - result = future.result() - if result: - final_result = {"url": future_to_url[future], "data": result} - break # Found a valid URL, stop checking others + for future in as_completed(future_to_url): + result = future.result() + if result: + return {"url": future_to_url[future], "data": result} - executor.shutdown(wait=False) # Don't wait for other threads to finish - return final_result + return None -def validate_account(): + +def validate_account() -> Tuple[Response, int]: + """Validates account credentials provided in a JSON request. + + Returns: + A Flask JSON response tuple containing a success or error message + and an HTTP status code. + """ data = request.get_json() username = data.get("username") password = data.get("password") @@ -54,6 +95,6 @@ def validate_account(): result = single_account_check(account_data, stream_urls) if result: - return jsonify({"message": "Account is valid"}) + return jsonify({"message": "Account is valid"}), 200 else: return jsonify({"message": "Account is invalid"}), 401 \ No newline at end of file diff --git a/ktvmanager/lib/database.py b/ktvmanager/lib/database.py index 647d2b3..8e2f678 100644 --- a/ktvmanager/lib/database.py +++ b/ktvmanager/lib/database.py @@ -1,12 +1,15 @@ import mysql.connector.pooling -from flask import jsonify, request, current_app +from flask import jsonify, request, current_app, Response from ktvmanager.lib.checker import single_account_check from ktvmanager.lib.encryption import encrypt_password, decrypt_password from ktvmanager.lib.get_urls import get_latest_urls_from_dns +from typing import List, Dict, Any, Optional, Tuple db_pool = None -def initialize_db_pool(): + +def initialize_db_pool() -> None: + """Initializes the database connection pool.""" global db_pool db_pool = mysql.connector.pooling.MySQLConnectionPool( pool_name="ktv_pool", @@ -15,10 +18,21 @@ def initialize_db_pool(): user=current_app.config["DBUSER"], password=current_app.config["DBPASS"], database=current_app.config["DATABASE"], - port=current_app.config["DBPORT"] + port=current_app.config["DBPORT"], ) -def _execute_query(query, params=None): + +def _execute_query(query: str, params: Optional[tuple] = None) -> List[Dict[str, Any]] | Dict[str, int]: + """Executes a SQL query and returns the result. + + Args: + query: The SQL query to execute. + params: The parameters to pass to the query. + + Returns: + A list of dictionaries for SELECT queries, or a dictionary with the + number of affected rows for other queries. + """ conn = db_pool.get_connection() cursor = conn.cursor(dictionary=True) try: @@ -33,31 +47,65 @@ def _execute_query(query, params=None): cursor.close() conn.close() -def get_user_id_from_username(username): + +def get_user_id_from_username(username: str) -> Optional[int]: + """Retrieves the user ID for a given username. + + Args: + username: The username to look up. + + Returns: + The user ID if found, otherwise None. + """ query = "SELECT id FROM users WHERE username = %s" result = _execute_query(query, (username,)) if result: - return result[0]['id'] + return result[0]["id"] return None -def get_user_accounts(user_id): + + +def get_user_accounts(user_id: int) -> Response: + """Retrieves all accounts for a given user ID. + + Args: + user_id: The ID of the user. + + Returns: + A Flask JSON response containing the user's accounts. + """ query = "SELECT * FROM userAccounts WHERE userID = %s" accounts = _execute_query(query, (user_id,)) for account in accounts: try: - account['password'] = decrypt_password(account['password']) + account["password"] = decrypt_password(account["password"]) except Exception as e: # Log the error to the console for debugging - print(f"Password decryption failed for account ID {account.get('id', 'N/A')}: {e}") - account['password'] = "DECRYPTION_FAILED" + print( + f"Password decryption failed for account ID {account.get('id', 'N/A')}: {e}" + ) + account["password"] = "DECRYPTION_FAILED" return jsonify(accounts) -def get_stream_names(): + +def get_stream_names() -> Response: + """Retrieves all stream names from the database. + + Returns: + A Flask JSON response containing a list of stream names. + """ query = "SELECT streamName FROM streams" results = _execute_query(query) - stream_names = [row['streamName'] for row in results] + stream_names = [row["streamName"] for row in results] return jsonify(stream_names) -def single_check(): + +def single_check() -> Response | Tuple[Response, int]: + """ + Performs a check on a single account provided in the request JSON. + + Returns: + A Flask JSON response with the result of the check, or an error message. + """ data = request.get_json() stream_urls = current_app.config["STREAM_URLS"] result = single_account_check(data, stream_urls) @@ -66,18 +114,43 @@ def single_check(): return jsonify(result) return jsonify({"message": "All checks failed"}), 400 -def add_account(user_id): + +def add_account(user_id: int) -> Response: + """Adds a new account for a user. + + Args: + user_id: The ID of the user. + + Returns: + A Flask JSON response confirming the account was added. + """ data = request.form res = single_account_check(data, get_latest_urls_from_dns()) - encrypted_password = encrypt_password(data['password']) + encrypted_password = encrypt_password(data["password"]) query = "INSERT INTO userAccounts (username, stream, streamURL, expiaryDate, password, userID) VALUES (%s, %s, %s, %s, %s, %s)" - params = (data['username'], data['stream'], res['url'], res['data']['user_info']['exp_date'], encrypted_password, user_id) + params = ( + data["username"], + data["stream"], + res["url"], + res["data"]["user_info"]["exp_date"], + encrypted_password, + user_id, + ) result = _execute_query(query, params) return jsonify(result) -def delete_account(user_id): + +def delete_account(user_id: int) -> Response: + """Deletes an account for a user. + + Args: + user_id: The ID of the user. + + Returns: + A Flask JSON response confirming the account was deleted. + """ data = request.form query = "DELETE FROM userAccounts WHERE username = %s AND stream = %s AND userId = %s" - params = (data['user'],data['stream'],user_id) + params = (data["user"], data["stream"], user_id) result = _execute_query(query, params) return jsonify(result) diff --git a/ktvmanager/lib/get_urls.py b/ktvmanager/lib/get_urls.py index 724137d..a2d0805 100644 --- a/ktvmanager/lib/get_urls.py +++ b/ktvmanager/lib/get_urls.py @@ -2,14 +2,20 @@ import os from dotenv import load_dotenv import requests import json +from typing import List load_dotenv() -def get_latest_urls_from_dns(): - with open('./ktvmanager/lib/DNS_list.txt') as f: - lines = [line.rstrip('\n') for line in f] - with open('./ktvmanager/lib/extra_urls.txt') as urls: - extra_urls = [line.rstrip('\n') for line in urls] +def get_latest_urls_from_dns() -> List[str]: + """Retrieves the latest stream URLs from DNS and extra URL files. + + Returns: + A list of unique stream URLs. + """ + with open("./ktvmanager/lib/DNS_list.txt") as f: + lines = [line.rstrip("\n") for line in f] + with open("./ktvmanager/lib/extra_urls.txt") as urls: + extra_urls = [line.rstrip("\n") for line in urls] # print(lines) complete_list_of_urls = [] @@ -21,20 +27,32 @@ def get_latest_urls_from_dns(): except Exception: pass try: - list_of_urls = content['su'].split(',') + list_of_urls = content["su"].split(",") except KeyError: - list_of_urls = content['fu'].split(',') + list_of_urls = content["fu"].split(",") for url in list_of_urls: complete_list_of_urls.append(url) complete_list_of_urls = list(set(complete_list_of_urls)) for url in extra_urls: complete_list_of_urls.append(url) return list(dict.fromkeys(complete_list_of_urls)) - -def generate_urls_for_user(username, password): + + +def generate_urls_for_user(username: str, password: str) -> List[str]: + """Generates a list of full stream URLs for a specific user. + + Args: + username: The username of the user. + password: The password of the user. + + Returns: + A list of fully constructed stream URLs for the user. + """ new_urls = [] for url in get_latest_urls_from_dns(): - hard_url = f'/player_api.php?password={password}&username={username}&action=user&sub=info' + hard_url = ( + f"/player_api.php?password={password}&username={username}&action=user&sub=info" + ) new_url = url + hard_url new_urls.append(new_url) return new_urls diff --git a/routes/api.py b/routes/api.py index 9b45633..8874c92 100644 --- a/routes/api.py +++ b/routes/api.py @@ -1,52 +1,140 @@ -from flask import Blueprint, jsonify -from ktvmanager.lib.database import get_user_accounts, get_stream_names, single_check, add_account, delete_account, get_user_id_from_username +from flask import Blueprint, jsonify, Response +from ktvmanager.lib.database import ( + get_user_accounts, + get_stream_names, + single_check, + add_account, + delete_account, + get_user_id_from_username, +) from ktvmanager.lib.get_urls import get_latest_urls_from_dns from ktvmanager.lib.auth import requires_basic_auth, check_login from ktvmanager.lib.checker import validate_account +from typing import Tuple api_blueprint = Blueprint("api", __name__) + @api_blueprint.route("/getUserAccounts") @requires_basic_auth -def get_user_accounts_route(username, password): +def get_user_accounts_route(username: str, password: str) -> Response: + """Retrieves all accounts associated with a user. + + Args: + username: The username of the user. + password: The password of the user (used for authentication). + + Returns: + A Flask JSON response containing the user's accounts or an error message. + """ user_id = get_user_id_from_username(username) if user_id: return get_user_accounts(user_id) return jsonify({"message": "User not found"}), 404 + @api_blueprint.route("/getStreamNames") @requires_basic_auth -def get_stream_names_route(username, password): +def get_stream_names_route(username: str, password: str) -> Response: + """Retrieves all stream names. + + Args: + username: The username of the user. + password: The password of the user (used for authentication). + + Returns: + A Flask JSON response containing the list of stream names. + """ return get_stream_names() + @api_blueprint.route("/getUserAccounts/streams") @requires_basic_auth -def get_user_accounts_streams_route(username, password): +def get_user_accounts_streams_route(username: str, password: str) -> Response: + """Retrieves the latest stream URLs from DNS. + + Args: + username: The username of the user. + password: The password of the user (used for authentication). + + Returns: + A Flask JSON response containing the list of stream URLs. + """ return jsonify(get_latest_urls_from_dns()) + @api_blueprint.route("/singleCheck", methods=["POST"]) @requires_basic_auth -def single_check_route(username, password): +def single_check_route(username: str, password: str) -> Response: + """Performs a single account check. + + Args: + username: The username of the user. + password: The password of the user (used for authentication). + + Returns: + A Flask JSON response with the result of the check. + """ return single_check() + @api_blueprint.route("/addAccount", methods=["POST"]) @requires_basic_auth -def add_account_route(username, password): +def add_account_route(username: str, password: str) -> Response: + """Adds a new account for the user. + + Args: + username: The username of the user. + password: The password of the user (used for authentication). + + Returns: + A Flask JSON response confirming the account was added or an error message. + """ user_id = get_user_id_from_username(username) return add_account(user_id) + @api_blueprint.route("/deleteAccount", methods=["POST"]) @requires_basic_auth -def delete_account_route(username, password): +def delete_account_route(username: str, password: str) -> Response: + """Deletes an account for the user. + + Args: + username: The username of the user. + password: The password of the user (used for authentication). + + Returns: + A Flask JSON response confirming the account was deleted or an error message. + """ user_id = get_user_id_from_username(username) return delete_account(user_id) + @api_blueprint.route("/validateAccount", methods=["POST"]) @requires_basic_auth -def validate_account_route(username, password): +def validate_account_route(username: str, password: str) -> Tuple[Response, int]: + """Validates an account. + + Args: + username: The username of the user. + password: The password of the user (used for authentication). + + Returns: + A tuple containing a Flask JSON response and an HTTP status code. + """ return validate_account() + @api_blueprint.route("/Login") @requires_basic_auth -def login_route(username, password): +def login_route(username: str, password: str) -> Response: + """Logs a user in. + + Args: + username: The username of the user. + password: The password of the user. + + Returns: + A Flask JSON response with the result of the login attempt. + """ return check_login(username, password) \ No newline at end of file