unshackle/unshackle/vaults/MySQL.py

from __future__ import annotations

import threading
from typing import Any, Iterator, Optional, Union
from uuid import UUID

import pymysql
from pymysql.cursors import DictCursor

from unshackle.core.config import config
from unshackle.core.vault import Vault, is_null_key


class MySQL(Vault):
    """Key Vault using a remotely-accessed mysql database connection."""

    def __init__(
        self, name: str, host: str, database: str, username: str, no_push: bool = False, **kwargs: Any
    ) -> None:
        """
        All extra arguments given through **kwargs will be sent to pymysql.connect.
        Use this to give more specific connection information.
        """
        super().__init__(name, no_push)
        self.slug = f"{host}:{database}:{username}"
        kwargs.setdefault("connect_timeout", config.vault_timeout or 10)
        self.conn_factory = ConnectionFactory(
            dict(host=host, db=database, user=username, cursorclass=DictCursor, **kwargs)
        )

        self.permissions = self.get_permissions()
        if not self.has_permission("SELECT"):
            raise PermissionError(f"MySQL vault {self.slug} has no SELECT permission.")

    def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
        if isinstance(kid, UUID):
            kid = kid.hex

        table = self.resolve_table(service)
        if not table:
            return None

        conn = self.conn_factory.get()
        cursor = conn.cursor()

        try:
            cursor.execute(
                # TODO: SQL injection risk
                f"SELECT `id`, `key_` FROM `{table}` WHERE `kid`=%s AND `key_`!=%s",
                (kid, "0" * 32),
            )
            cek = cursor.fetchone()
            return cek["key_"] if cek else None
        finally:
            cursor.close()

    def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
        table = self.resolve_table(service)
        if not table:
            return None

        conn = self.conn_factory.get()
        cursor = conn.cursor()

        try:
            cursor.execute(
                # TODO: SQL injection risk
                f"SELECT `kid`, `key_` FROM `{table}` WHERE `key_`!=%s",
                ("0" * 32,),
            )
            for row in cursor.fetchall():
                yield row["kid"], row["key_"]
        finally:
            cursor.close()

    def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
        if is_null_key(key):
            raise ValueError("You cannot add a NULL Content Key to a Vault.")

        service = self.resolve_table(service) or service

        if not self.has_permission("INSERT", table=service):
            raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.")

        if not self.has_table(service):
            try:
                self.create_table(service)
            except PermissionError:
                return False

        if isinstance(kid, UUID):
            kid = kid.hex

        conn = self.conn_factory.get()
        cursor = conn.cursor()

        try:
            cursor.execute(
                # TODO: SQL injection risk
                f"SELECT `id` FROM `{service}` WHERE `kid`=%s AND `key_`=%s",
                (kid, key),
            )
            if cursor.fetchone():
                return True
            cursor.execute(
                # TODO: SQL injection risk
                f"INSERT INTO `{service}` (kid, key_) VALUES (%s, %s)",
                (kid, key),
            )
        finally:
            conn.commit()
            cursor.close()

        return True

    def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
        for kid, key in kid_keys.items():
            if is_null_key(key):
                raise ValueError("You cannot add a NULL Content Key to a Vault.")

        service = self.resolve_table(service) or service

        if not self.has_permission("INSERT", table=service):
            raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.")

        if not self.has_table(service):
            try:
                self.create_table(service)
            except PermissionError:
                return 0

        if not isinstance(kid_keys, dict):
            raise ValueError(f"The kid_keys provided is not a dictionary, {kid_keys!r}")
        if not all(isinstance(kid, (str, UUID)) and isinstance(key_, str) for kid, key_ in kid_keys.items()):
            raise ValueError("Expecting dict with Key of str/UUID and value of str.")

        if any(isinstance(kid, UUID) for kid, key_ in kid_keys.items()):
            kid_keys = {kid.hex if isinstance(kid, UUID) else kid: key_ for kid, key_ in kid_keys.items()}

        if not kid_keys:
            return 0

        conn = self.conn_factory.get()
        cursor = conn.cursor()

        try:
            existing_kids: set[str] = set()
            kid_list = list(kid_keys)
            for i in range(0, len(kid_list), 10000):
                batch = kid_list[i : i + 10000]
                placeholders = ",".join(["%s"] * len(batch))
                cursor.execute(f"SELECT kid FROM `{service}` WHERE kid IN ({placeholders})", batch)
                existing_kids.update(row["kid"] for row in cursor.fetchall())

            new_keys = {kid: key for kid, key in kid_keys.items() if kid not in existing_kids}

            if not new_keys:
                return 0

            cursor.executemany(
                f"INSERT INTO `{service}` (kid, key_) VALUES (%s, %s)",
                new_keys.items(),
            )
            return len(new_keys)
        finally:
            conn.commit()
            cursor.close()

    def get_services(self) -> Iterator[str]:
        conn = self.conn_factory.get()
        cursor = conn.cursor()

        try:
            cursor.execute("SHOW TABLES")
            for table in cursor.fetchall():
                # each entry has a key named `Tables_in_<db name>`
                yield list(table.values())[0]
        finally:
            cursor.close()

    def resolve_table(self, name: str) -> Optional[str]:
        """
        Get the actual Table name matching `name` case-insensitively, if it exists.

        The UPPER() comparison is explicit so the result does not depend on the server's
        `lower_case_table_names` setting. Where both spellings exist, exact case wins.
        """
        conn = self.conn_factory.get()
        cursor = conn.cursor()

        try:
            cursor.execute(
                "SELECT TABLE_NAME FROM information_schema.TABLES"
                " WHERE TABLE_SCHEMA=%s AND UPPER(TABLE_NAME)=UPPER(%s)",
                (conn.db, name),
            )
            names = [list(row.values())[0] for row in cursor.fetchall()]
        finally:
            cursor.close()

        if name in names:
            return name
        return names[0] if names else None

    def has_table(self, name: str) -> bool:
        """Return True if the Vault has a Table with the specified name."""
        conn = self.conn_factory.get()
        cursor = conn.cursor()

        try:
            cursor.execute(
                "SELECT count(TABLE_NAME) FROM information_schema.TABLES WHERE TABLE_SCHEMA=%s AND TABLE_NAME=%s",
                (conn.db, name),
            )
            row = cursor.fetchone()
            assert row is not None  # COUNT() always returns one row
            return list(row.values())[0] == 1
        finally:
            cursor.close()

    def create_table(self, name: str) -> None:
        """Make a Table with the specified name if it does not exist yet."""
        if self.has_table(name):
            return

        if not self.has_permission("CREATE"):
            raise PermissionError(f"MySQL vault {self.slug} has no CREATE permission.")

        conn = self.conn_factory.get()
        cursor = conn.cursor()

        try:
            cursor.execute(
                # TODO: SQL injection risk
                f"""
                CREATE TABLE IF NOT EXISTS {name} (
                  id          int AUTO_INCREMENT PRIMARY KEY,
                  kid         VARCHAR(64) NOT NULL,
                  key_        VARCHAR(64) NOT NULL,
                  UNIQUE(kid, key_)
                );
                """
            )
        finally:
            conn.commit()
            cursor.close()

    def get_permissions(self) -> list:
        """Get and parse Grants to a more usable list tuple array."""
        conn = self.conn_factory.get()
        cursor = conn.cursor()

        try:
            cursor.execute("SHOW GRANTS")
            grants = [next(iter(x.values())) for x in cursor.fetchall()]
            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
        finally:
            conn.commit()
            cursor.close()

    def has_permission(self, operation: str, database: Optional[str] = None, table: Optional[str] = None) -> bool:
        """Return True if the current connection has a specific permission."""
        grants = [x for x in self.permissions 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 ConnectionFactory:
    def __init__(self, con: dict) -> None:
        self._con = con
        self._store = threading.local()

    def create_connection(self) -> pymysql.Connection[DictCursor]:
        return pymysql.connect(**self._con)

    def get(self) -> pymysql.Connection[DictCursor]:
        if not hasattr(self._store, "conn"):
            self._store.conn = self.create_connection()
        return self._store.conn