#!/usr/bin/env python3 import argparse import json import sqlite3 from vinetrimmer.utils.AtomicSQL import AtomicSQL parser = argparse.ArgumentParser( "Key Store DB merger", description="Simple script to merge vinetrimmer key store db's into one" ) parser.add_argument( "-i", "--input", help="key store db that will send keys", required=True) parser.add_argument( "-o", "--output", help="key store db that will receive keys", required=True) args = parser.parse_args() add_count = 0 update_count = 0 existed_count = 0 input_db = AtomicSQL() input_db_id = input_db.load(sqlite3.connect(args.input)) output_db = AtomicSQL() output_db_id = output_db.load(sqlite3.connect(args.output)) # get all keys from input db input_keys = input_db.safe_execute( input_db_id, lambda db, cursor: cursor.execute("SELECT * FROM `keys`") ).fetchall() for i, service, title, pssh_b64, pssh_sha1, content_keys in input_keys: exists = output_db.safe_execute( output_db_id, lambda db, cursor: cursor.execute( """ SELECT "id","service","title","pssh_b64","pssh_sha1","content_keys" FROM `keys` WHERE `service`=:service AND (`pssh_b64`=:pssh_b64 or `pssh_sha1`=:pssh_sha1) """, { "service": service, "pssh_b64": pssh_b64, "pssh_sha1": pssh_sha1 } ) ).fetchone() if exists: has_differences = ( json.loads(exists[5]) != json.loads(content_keys) or title != exists[2] or pssh_b64 != exists[3] or pssh_sha1 != exists[4] ) if has_differences: update_count += 1 content_keys = list(set(json.loads(exists[5])) | set(json.loads(content_keys))) print(f"Updating {title} {service} {pssh_b64}: {content_keys}") output_db.safe_execute( output_db_id, lambda db, cursor: cursor.execute( """ UPDATE `keys` SET `service`=:service, `title`=:title, `pssh_b64`=:new_pssh_b64, `pssh_sha1`=:new_pssh_sha1, `content_keys`=:content_keys WHERE `service`=:service AND (`pssh_b64`=:pssh_b64 or `pssh_sha1`=:pssh_sha1) """, { "service": service, "title": title or exists[2], "pssh_b64": pssh_b64, "new_pssh_b64": pssh_b64 or exists[3], "pssh_sha1": pssh_sha1, "new_pssh_sha1": pssh_sha1 or exists[4], "content_keys": json.dumps(content_keys, separators=(",", ":")) } ) ) else: existed_count += 1 print(f"Key {title} {service} {pssh_b64} already exists in the db with no differences, skipping...") else: add_count += 1 print(f"Adding {title} {service} {pssh_b64}: {content_keys}") output_db.safe_execute( output_db_id, lambda db, cursor: cursor.execute( """ INSERT INTO `keys` (service, title, pssh_b64, pssh_sha1, content_keys) VALUES (:service, :title, :pssh_b64, :pssh_sha1, :content_keys) """, { "service": service, "title": title, "pssh_b64": pssh_b64, "pssh_sha1": pssh_sha1, "content_keys": json.dumps(content_keys, separators=(",", ":")) } ) ) output_db.commit(output_db_id) print( "Done!\n" f"{add_count} added, {update_count} updated in some way, {existed_count} already existed (no difference)" )