|
1 | 1 | #!/usr/bin/env python3 |
2 | 2 |
|
3 | | -import datetime |
| 3 | +from datetime import datetime |
4 | 4 | import os |
5 | 5 | import logging |
6 | 6 |
|
| 7 | +#-------------------------------------------------------------------------------------------------------------------------------------------------------- |
| 8 | + |
| 9 | +loggersrv = logging.getLogger('logsrv') |
| 10 | +_column_names = ('clientMachineId', 'machineName', 'applicationId', 'skuId', 'licenseStatus', 'lastRequestTime', 'kmsEpid', 'requestCount', 'lastRequestIP') |
| 11 | + |
7 | 12 | # sqlite3 is optional. |
| 13 | +available = False |
8 | 14 | try: |
9 | 15 | import sqlite3 |
| 16 | + available = True |
10 | 17 | except ImportError: |
11 | 18 | pass |
12 | 19 |
|
13 | | -from pykms_Format import pretty_printer |
14 | | - |
15 | | -#-------------------------------------------------------------------------------------------------------------------------------------------------------- |
16 | | - |
17 | | -loggersrv = logging.getLogger('logsrv') |
18 | | - |
19 | 20 | def sql_initialize(dbName): |
| 21 | + if available is False: |
| 22 | + loggersrv.info("'sqlite3' module not found! SQLite database support cannot be enabled.") |
| 23 | + return |
| 24 | + loggersrv.debug(f'SQLite database support enabled. Database file: "{dbName}"') |
20 | 25 | if not os.path.isfile(dbName): |
21 | | - # Initialize the database. |
| 26 | + # Initialize the database |
22 | 27 | loggersrv.debug(f'Initializing database file "{dbName}"...') |
23 | | - con = None |
24 | | - try: |
25 | | - con = sqlite3.connect(dbName) |
| 28 | + with sqlite3.connect(dbName) as con: |
26 | 29 | cur = con.cursor() |
27 | | - cur.execute("CREATE TABLE clients(clientMachineId TEXT , machineName TEXT, applicationId TEXT, skuId TEXT, licenseStatus TEXT, lastRequestTime INTEGER, kmsEpid TEXT, requestCount INTEGER, PRIMARY KEY(clientMachineId, applicationId))") |
| 30 | + cur.execute("CREATE TABLE clients(clientMachineId TEXT, machineName TEXT, applicationId TEXT, skuId TEXT, licenseStatus TEXT, lastRequestTime INTEGER, kmsEpid TEXT, requestCount INTEGER, PRIMARY KEY(clientMachineId, applicationId))") |
| 31 | + |
| 32 | + if os.path.isfile(dbName): |
| 33 | + # Update database |
| 34 | + with sqlite3.connect(dbName) as con: |
| 35 | + cur = con.cursor() |
| 36 | + # Create simple "metadata" table if not exists. |
| 37 | + cur.execute("CREATE TABLE IF NOT EXISTS metadata (key TEXT PRIMARY KEY, value TEXT);") |
| 38 | + # Get the current schema version |
| 39 | + cur.execute("SELECT value FROM metadata WHERE key='schema_version';") |
| 40 | + row = cur.fetchone() |
| 41 | + if row is None: |
| 42 | + current_version = 0 |
| 43 | + else: |
| 44 | + current_version = int(row[0]) |
| 45 | + loggersrv.debug(f'Current database schema version: {current_version}') |
| 46 | + # Apply necessary migrations |
| 47 | + if current_version < 1: |
| 48 | + # v1: Add "lastRequestIP" column to "clients" table. |
| 49 | + loggersrv.info("Upgrading database schema to version 1...") |
| 50 | + cur.execute("ALTER TABLE clients ADD COLUMN lastRequestIP TEXT;") |
| 51 | + cur.execute("INSERT OR REPLACE INTO metadata (key, value) VALUES ('schema_version', '1');") |
| 52 | + loggersrv.info("Database schema updated to version 1.") |
28 | 53 |
|
29 | | - except sqlite3.Error as e: |
30 | | - pretty_printer(log_obj = loggersrv.error, to_exit = True, put_text = "{reverse}{red}{bold}Sqlite Error: %s. Exiting...{end}" %str(e)) |
31 | | - finally: |
32 | | - if con: |
33 | | - con.commit() |
34 | | - con.close() |
35 | 54 |
|
36 | 55 | def sql_get_all(dbName): |
| 56 | + if available is False: |
| 57 | + return |
37 | 58 | if not os.path.isfile(dbName): |
38 | 59 | return None |
39 | 60 | with sqlite3.connect(dbName) as con: |
| 61 | + con.row_factory = sqlite3.Row |
40 | 62 | cur = con.cursor() |
41 | | - cur.execute("SELECT * FROM clients") |
| 63 | + cur.execute(f"SELECT {', '.join(_column_names)} FROM clients") |
42 | 64 | clients = [] |
43 | 65 | for row in cur.fetchall(): |
44 | | - clients.append({ |
45 | | - 'clientMachineId': row[0], |
46 | | - 'machineName': row[1], |
47 | | - 'applicationId': row[2], |
48 | | - 'skuId': row[3], |
49 | | - 'licenseStatus': row[4], |
50 | | - 'lastRequestTime': datetime.datetime.fromtimestamp(row[5]).isoformat(), |
51 | | - 'kmsEpid': row[6], |
52 | | - 'requestCount': row[7] |
53 | | - }) |
| 66 | + loggersrv.debug(f"Row: {row}") |
| 67 | + obj = {} |
| 68 | + for col_name in _column_names: |
| 69 | + if col_name == "lastRequestTime": |
| 70 | + obj[col_name] = datetime.fromtimestamp(row['lastRequestTime']).isoformat() |
| 71 | + else: |
| 72 | + obj[col_name] = row[col_name] |
| 73 | + loggersrv.debug(f"Obj: {obj}") |
| 74 | + clients.append(obj) |
54 | 75 | return clients |
55 | 76 |
|
56 | 77 | def sql_update(dbName, infoDict): |
57 | | - con = None |
58 | | - try: |
59 | | - con = sqlite3.connect(dbName) |
| 78 | + if available is False: |
| 79 | + return |
| 80 | + |
| 81 | + # make sure all column names are present |
| 82 | + for col_name in _column_names: |
| 83 | + if col_name in ["requestCount", "kmsEpid"]: |
| 84 | + continue |
| 85 | + if col_name not in infoDict: |
| 86 | + raise ValueError(f"infoDict is missing required column: {col_name}") |
| 87 | + |
| 88 | + with sqlite3.connect(dbName) as con: |
| 89 | + con.row_factory = sqlite3.Row |
60 | 90 | cur = con.cursor() |
61 | | - cur.execute("SELECT * FROM clients WHERE clientMachineId=:clientMachineId AND applicationId=:appId;", infoDict) |
62 | | - try: |
63 | | - data = cur.fetchone() |
64 | | - if not data: |
65 | | - # Insert row. |
66 | | - cur.execute("INSERT INTO clients (clientMachineId, machineName, applicationId, \ |
67 | | -skuId, licenseStatus, lastRequestTime, requestCount) VALUES (:clientMachineId, :machineName, :appId, :skuId, :licenseStatus, :requestTime, 1);", infoDict) |
68 | | - else: |
69 | | - # Update data. |
70 | | - if data[1] != infoDict["machineName"]: |
71 | | - cur.execute("UPDATE clients SET machineName=:machineName WHERE \ |
72 | | -clientMachineId=:clientMachineId AND applicationId=:appId;", infoDict) |
73 | | - if data[2] != infoDict["appId"]: |
74 | | - cur.execute("UPDATE clients SET applicationId=:appId WHERE \ |
75 | | -clientMachineId=:clientMachineId AND applicationId=:appId;", infoDict) |
76 | | - if data[3] != infoDict["skuId"]: |
77 | | - cur.execute("UPDATE clients SET skuId=:skuId WHERE \ |
78 | | -clientMachineId=:clientMachineId AND applicationId=:appId;", infoDict) |
79 | | - if data[4] != infoDict["licenseStatus"]: |
80 | | - cur.execute("UPDATE clients SET licenseStatus=:licenseStatus WHERE \ |
81 | | -clientMachineId=:clientMachineId AND applicationId=:appId;", infoDict) |
82 | | - if data[5] != infoDict["requestTime"]: |
83 | | - cur.execute("UPDATE clients SET lastRequestTime=:requestTime WHERE \ |
84 | | -clientMachineId=:clientMachineId AND applicationId=:appId;", infoDict) |
85 | | - # Increment requestCount |
86 | | - cur.execute("UPDATE clients SET requestCount=requestCount+1 WHERE \ |
87 | | -clientMachineId=:clientMachineId AND applicationId=:appId;", infoDict) |
88 | | - |
89 | | - except sqlite3.Error as e: |
90 | | - pretty_printer(log_obj = loggersrv.error, to_exit = True, |
91 | | - put_text = "{reverse}{red}{bold}Sqlite Error: %s. Exiting...{end}" %str(e)) |
92 | | - except sqlite3.Error as e: |
93 | | - pretty_printer(log_obj = loggersrv.error, to_exit = True, |
94 | | - put_text = "{reverse}{red}{bold}Sqlite Error: %s. Exiting...{end}" %str(e)) |
95 | | - finally: |
96 | | - if con: |
97 | | - con.commit() |
98 | | - con.close() |
| 91 | + cur.execute(f"SELECT {', '.join(_column_names)} FROM clients WHERE clientMachineId=:clientMachineId AND applicationId=:applicationId;", infoDict) |
| 92 | + data = cur.fetchone() |
| 93 | + if not data: |
| 94 | + # Insert new row with all given info |
| 95 | + infoDict["kmsEpid"] = "" # Default empty value |
| 96 | + infoDict["requestCount"] = 1 |
| 97 | + cur.execute(f"""INSERT INTO clients ({', '.join(_column_names)}) |
| 98 | + VALUES ({', '.join(':' + col for col in _column_names)});""", infoDict) |
| 99 | + |
| 100 | + else: |
| 101 | + # Update only changed columns |
| 102 | + common_postfix = "WHERE clientMachineId=:clientMachineId AND applicationId=:applicationId" |
| 103 | + def update_column_if_changed(column_name, new_value): |
| 104 | + assert "clientMachineId" in infoDict and "applicationId" in infoDict, "infoDict must contain 'clientMachineId' and 'applicationId'" |
| 105 | + if column_name not in _column_names: |
| 106 | + raise ValueError(f"Unknown column name: {column_name}") |
| 107 | + if data[column_name] != new_value: |
| 108 | + query = f"UPDATE clients SET {column_name}=:value {common_postfix}" |
| 109 | + cur.execute(query, {"value": new_value, "clientMachineId": infoDict['clientMachineId'], "applicationId": infoDict['applicationId']}) |
| 110 | + |
| 111 | + # Dynamically check and maybe update all columns |
| 112 | + for column_name in _column_names: |
| 113 | + if column_name in ["clientMachineId", "applicationId", "requestCount"]: |
| 114 | + continue # Skip these columns |
| 115 | + if column_name == "kmsEpid": |
| 116 | + # this one can only be updated by the special function |
| 117 | + continue |
| 118 | + update_column_if_changed(column_name, infoDict[column_name]) |
| 119 | + |
| 120 | + # Finally increment requestCount |
| 121 | + cur.execute(f"UPDATE clients SET requestCount=requestCount+1 {common_postfix}", infoDict) |
99 | 122 |
|
100 | 123 | def sql_update_epid(dbName, kmsRequest, response, appName): |
| 124 | + if available is False: |
| 125 | + return |
| 126 | + |
101 | 127 | cmid = str(kmsRequest['clientMachineId'].get()) |
102 | | - con = None |
103 | | - try: |
104 | | - con = sqlite3.connect(dbName) |
| 128 | + with sqlite3.connect(dbName) as con: |
105 | 129 | cur = con.cursor() |
106 | | - cur.execute("SELECT * FROM clients WHERE clientMachineId=? AND applicationId=?;", (cmid, appName)) |
107 | | - try: |
108 | | - data = cur.fetchone() |
109 | | - cur.execute("UPDATE clients SET kmsEpid=? WHERE \ |
110 | | -clientMachineId=? AND applicationId=?;", (str(response["kmsEpid"].decode('utf-16le')), cmid, appName)) |
111 | | - |
112 | | - except sqlite3.Error as e: |
113 | | - pretty_printer(log_obj = loggersrv.error, to_exit = True, |
114 | | - put_text = "{reverse}{red}{bold}Sqlite Error: %s. Exiting...{end}" %str(e)) |
115 | | - except sqlite3.Error as e: |
116 | | - pretty_printer(log_obj = loggersrv.error, to_exit = True, |
117 | | - put_text = "{reverse}{red}{bold}Sqlite Error: %s. Exiting...{end}" %str(e)) |
118 | | - finally: |
119 | | - if con: |
120 | | - con.commit() |
121 | | - con.close() |
| 130 | + cur.execute("UPDATE clients SET kmsEpid=? WHERE clientMachineId=? AND applicationId=?;", |
| 131 | + (str(response["kmsEpid"].decode('utf-16le')), cmid, appName)) |
0 commit comments