V2.1.0 upstream

vaults/MySQL.py

@@ -28,26 +28,33 @@ class MySQL(Vault):
             raise PermissionError(f"MySQL vault {self.slug} has no SELECT permission.")
 
     def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
-        if not self.has_table(service):
-            # no table, no key, simple
-            return None
-
         if isinstance(kid, UUID):
             kid = kid.hex
 
+        service_variants = [service]
+        if service != service.lower():
+            service_variants.append(service.lower())
+        if service != service.upper():
+            service_variants.append(service.upper())
+
         conn = self.conn_factory.get()
         cursor = conn.cursor()
 
         try:
-            cursor.execute(
-                # TODO: SQL injection risk
-                f"SELECT `id`, `key_` FROM `{service}` WHERE `kid`=%s AND `key_`!=%s",
-                (kid, "0" * 32),
-            )
-            cek = cursor.fetchone()
-            if not cek:
-                return None
-            return cek["key_"]
+            for service_name in service_variants:
+                if not self.has_table(service_name):
+                    continue
+
+                cursor.execute(
+                    # TODO: SQL injection risk
+                    f"SELECT `id`, `key_` FROM `{service_name}` WHERE `kid`=%s AND `key_`!=%s",
+                    (kid, "0" * 32),
+                )
+                cek = cursor.fetchone()
+                if cek:
+                    return cek["key_"]
+
+            return None
         finally:
             cursor.close()
 
@@ -131,16 +138,27 @@ class MySQL(Vault):
         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:
+            placeholders = ",".join(["%s"] * len(kid_keys))
+            cursor.execute(f"SELECT kid FROM `{service}` WHERE kid IN ({placeholders})", list(kid_keys.keys()))
+            existing_kids = {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(
-                # TODO: SQL injection risk
-                f"INSERT IGNORE INTO `{service}` (kid, key_) VALUES (%s, %s)",
-                kid_keys.items(),
+                f"INSERT INTO `{service}` (kid, key_) VALUES (%s, %s)",
+                new_keys.items(),
             )
-            return cursor.rowcount
+            return len(new_keys)
         finally:
             conn.commit()
             cursor.close()