223 lines
8.8 KiB
Python
223 lines
8.8 KiB
Python
import os
|
|
import sqlite3
|
|
from enum import Enum
|
|
|
|
import pymysql
|
|
|
|
from vinetrimmer.utils.AtomicSQL import AtomicSQL
|
|
|
|
|
|
class InsertResult(Enum):
|
|
FAILURE = 0
|
|
SUCCESS = 1
|
|
ALREADY_EXISTS = 2
|
|
|
|
|
|
class Vault:
|
|
"""
|
|
Key Vault.
|
|
This defines various details about the vault, including its Connection object.
|
|
"""
|
|
|
|
def __init__(self, type_, name, ticket=None, path=None, username=None, password=None, database=None,
|
|
host=None, port=3306):
|
|
from vinetrimmer.config import directories
|
|
|
|
try:
|
|
self.type = self.Types[type_.upper()]
|
|
except KeyError:
|
|
raise ValueError(f"Invalid vault type [{type_}]")
|
|
self.name = name
|
|
self.con = None
|
|
if self.type == Vault.Types.LOCAL:
|
|
if not path:
|
|
raise ValueError("Local vault has no path specified")
|
|
self.con = sqlite3.connect(os.path.expanduser(path).format(data_dir=directories.data))
|
|
elif self.type == Vault.Types.REMOTE:
|
|
self.con = pymysql.connect(
|
|
user=username,
|
|
password=password or "",
|
|
db=database,
|
|
host=host,
|
|
port=port,
|
|
cursorclass=pymysql.cursors.DictCursor # TODO: Needed? Maybe use it on sqlite3 too?
|
|
)
|
|
else:
|
|
raise ValueError(f"Invalid vault type [{self.type.name}]")
|
|
self.ph = {self.Types.LOCAL: "?", self.Types.REMOTE: "%s"}[self.type]
|
|
self.ticket = ticket
|
|
|
|
self.perms = self.get_permissions()
|
|
if not self.has_permission("SELECT"):
|
|
raise ValueError(f"Cannot use vault. Vault {self.name} has no SELECT permission.")
|
|
|
|
def __str__(self):
|
|
return f"{self.name} ({self.type.name})"
|
|
|
|
def get_permissions(self):
|
|
if self.type == self.Types.LOCAL:
|
|
return [tuple([["*"], tuple(["*", "*"])])]
|
|
|
|
with self.con.cursor() as c:
|
|
c.execute("SHOW GRANTS")
|
|
grants = c.fetchall()
|
|
grants = [next(iter(x.values())) for x in grants]
|
|
grants = [tuple(x[6:].split(" TO ")[0].split(" ON ")) for x in list(grants)]
|
|
grants = [(
|
|
list(map(str.strip, perms.replace("ALL PRIVILEGES", "*").split(","))),
|
|
location.replace("`", "").split(".")
|
|
) for perms, location in grants]
|
|
|
|
return grants
|
|
|
|
def has_permission(self, operation, database=None, table=None):
|
|
grants = [x for x in self.perms if x[0] == ["*"] or operation.upper() in x[0]]
|
|
if grants and database:
|
|
grants = [x for x in grants if x[1][0] in (database, "*")]
|
|
if grants and table:
|
|
grants = [x for x in grants if x[1][1] in (table, "*")]
|
|
return bool(grants)
|
|
|
|
class Types(Enum):
|
|
LOCAL = 1
|
|
REMOTE = 2
|
|
|
|
|
|
class Vaults:
|
|
"""
|
|
Key Vaults.
|
|
Keeps hold of Vault objects, with convenience functions for
|
|
using multiple vaults in one actions, e.g. searching vaults
|
|
for a key based on kid.
|
|
This object uses AtomicSQL for accessing the vault connections
|
|
instead of directly. This is to provide thread safety but isn't
|
|
strictly necessary.
|
|
"""
|
|
|
|
def __init__(self, vaults, service):
|
|
self.adb = AtomicSQL()
|
|
self.vaults = sorted(vaults, key=lambda v: 0 if v.type == Vault.Types.LOCAL else 1)
|
|
self.service = service.lower()
|
|
for vault in self.vaults:
|
|
vault.ticket = self.adb.load(vault.con)
|
|
self.create_table(vault, self.service, commit=True)
|
|
|
|
def __iter__(self):
|
|
return iter(self.vaults)
|
|
|
|
def get(self, kid, title):
|
|
for vault in self.vaults:
|
|
# Note on why it matches by KID instead of PSSH:
|
|
# Matching cache by pssh is not efficient. The PSSH can be made differently by all different
|
|
# clients for all different reasons, e.g. only having the init data, but the cached PSSH is
|
|
# a manually crafted PSSH, which may not match other clients manually crafted PSSH, and such.
|
|
# So it searches by KID instead for this reason, as the KID has no possibility of being different
|
|
# client to client other than capitalization. There is an unknown with KID matching, It's unknown
|
|
# for *sure* if the KIDs ever conflict or not with another bitrate/stream/title. I haven't seen
|
|
# this happen ever and neither has anyone I have asked.
|
|
if not self.table_exists(vault, self.service):
|
|
continue # this service has no service table, so no keys, just skip
|
|
if not vault.ticket:
|
|
raise ValueError(f"Vault {vault.name} does not have a valid ticket available.")
|
|
c = self.adb.safe_execute(
|
|
vault.ticket,
|
|
lambda db, cursor: cursor.execute(
|
|
"SELECT `id`, `key_`, `title` FROM `{1}` WHERE `kid`={0}".format(vault.ph, self.service),
|
|
[kid]
|
|
)
|
|
).fetchone()
|
|
if c:
|
|
if isinstance(c, dict):
|
|
c = list(c.values())
|
|
if not c[2] and vault.has_permission("UPDATE", table=self.service):
|
|
self.adb.safe_execute(
|
|
vault.ticket,
|
|
lambda db, cursor: cursor.execute(
|
|
"UPDATE `{1}` SET `title`={0} WHERE `id`={0}".format(vault.ph, self.service),
|
|
[title, c[0]]
|
|
)
|
|
)
|
|
self.commit(vault)
|
|
return c[1], vault
|
|
return None, None
|
|
|
|
def table_exists(self, vault, table):
|
|
if not vault.ticket:
|
|
raise ValueError(f"Vault {vault.name} does not have a valid ticket available.")
|
|
if vault.type == Vault.Types.LOCAL:
|
|
return self.adb.safe_execute(
|
|
vault.ticket,
|
|
lambda db, cursor: cursor.execute(
|
|
f"SELECT count(name) FROM sqlite_master WHERE type='table' AND name={vault.ph}",
|
|
[table]
|
|
)
|
|
).fetchone()[0] == 1
|
|
return list(self.adb.safe_execute(
|
|
vault.ticket,
|
|
lambda db, cursor: cursor.execute(
|
|
f"SELECT count(TABLE_NAME) FROM information_schema.TABLES WHERE TABLE_NAME={vault.ph}",
|
|
[table]
|
|
)
|
|
).fetchone().values())[0] == 1
|
|
|
|
def create_table(self, vault, table, commit=False):
|
|
if self.table_exists(vault, table):
|
|
return
|
|
if not vault.ticket:
|
|
raise ValueError(f"Vault {vault.name} does not have a valid ticket available.")
|
|
if vault.has_permission("CREATE"):
|
|
print(f"Creating `{table}` table in {vault} key vault...")
|
|
self.adb.safe_execute(
|
|
vault.ticket,
|
|
lambda db, cursor: cursor.execute(
|
|
"CREATE TABLE IF NOT EXISTS {} (".format(table) + (
|
|
"""
|
|
"id" INTEGER NOT NULL UNIQUE,
|
|
"kid" TEXT NOT NULL COLLATE NOCASE,
|
|
"key_" TEXT NOT NULL COLLATE NOCASE,
|
|
"title" TEXT,
|
|
PRIMARY KEY("id" AUTOINCREMENT),
|
|
UNIQUE("kid", "key_")
|
|
""" if vault.type == Vault.Types.LOCAL else
|
|
"""
|
|
id INTEGER AUTO_INCREMENT PRIMARY KEY,
|
|
kid VARCHAR(255) NOT NULL,
|
|
key_ VARCHAR(255) NOT NULL,
|
|
title TEXT,
|
|
UNIQUE(kid, key_)
|
|
"""
|
|
) + ");"
|
|
)
|
|
)
|
|
if commit:
|
|
self.commit(vault)
|
|
|
|
def insert_key(self, vault, table, kid, key, title, commit=False):
|
|
if not self.table_exists(vault, table):
|
|
return InsertResult.FAILURE
|
|
if not vault.ticket:
|
|
raise ValueError(f"Vault {vault.name} does not have a valid ticket available.")
|
|
if not vault.has_permission("INSERT", table=table):
|
|
raise ValueError(f"Cannot insert key into Vault. Vault {vault.name} has no INSERT permission.")
|
|
if self.adb.safe_execute(
|
|
vault.ticket,
|
|
lambda db, cursor: cursor.execute(
|
|
"SELECT `id` FROM `{1}` WHERE `kid`={0} AND `key_`={0}".format(vault.ph, self.service),
|
|
[kid, key]
|
|
)
|
|
).fetchone():
|
|
return InsertResult.ALREADY_EXISTS
|
|
self.adb.safe_execute(
|
|
vault.ticket,
|
|
lambda db, cursor: cursor.execute(
|
|
"INSERT INTO `{1}` (kid, key_, title) VALUES ({0}, {0}, {0})".format(vault.ph, table),
|
|
(kid, key, title)
|
|
)
|
|
)
|
|
if commit:
|
|
self.commit(vault)
|
|
return InsertResult.SUCCESS
|
|
|
|
def commit(self, vault):
|
|
self.adb.commit(vault.ticket)
|