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)