docstrings

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 = {
+    account_data = {"username": args.username, "password": args.password}
-        "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()

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'],
+            "username": account["username"],
-            "password": decrypted_password
+            "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()

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 (
-        return f(auth.username,auth.password, *args, **kwargs)
+                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"})

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)))
+    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'])
+                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}
+                return {"url": future_to_url[future], "data": result}
-            break  # Found a valid URL, stop checking others
 
-    executor.shutdown(wait=False)  # Don't wait for other threads to finish
+    return None
-    return final_result
 
-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

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}")
+            print(
-            account['password'] = "DECRYPTION_FAILED"
+                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)

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():
+def get_latest_urls_from_dns() -> List[str]:
-    with open('./ktvmanager/lib/DNS_list.txt') as f:
+    """Retrieves the latest stream URLs from DNS and extra URL files.
-        lines = [line.rstrip('\n') for line in f]
+
-    with open('./ktvmanager/lib/extra_urls.txt') as urls:
+    Returns:
-        extra_urls = [line.rstrip('\n') for line in urls]        
+        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,9 +27,9 @@ 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))
@@ -31,10 +37,22 @@ def get_latest_urls_from_dns():
         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

routes/api.py

@@ -1,52 +1,140 @@
-from flask import Blueprint, jsonify
+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.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)