113 lines
3.8 KiB
Python
113 lines
3.8 KiB
Python
#!/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)"
|
|
)
|