101 lines
2.9 KiB
Python
101 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import re
|
|
import sqlite3
|
|
import sys
|
|
|
|
from vinetrimmer.utils.AtomicSQL import AtomicSQL
|
|
|
|
"""
|
|
Add keys to key vault. File should have one KID:KEY per-line.
|
|
Optionally you can also put `:<title here>` at the end (after `KEY`).
|
|
"""
|
|
|
|
parser = argparse.ArgumentParser(
|
|
"Key Vault DB batch adder/updater",
|
|
description="Simple script to add or update key information to a vinetrimmer key vault db"
|
|
)
|
|
parser.add_argument(
|
|
"-t", "--table",
|
|
help="table to store keys to. (e.g. amazon, netflix, disneyplus)",
|
|
required=True)
|
|
parser.add_argument(
|
|
"-i", "--input",
|
|
help="data used to parse from",
|
|
required=True)
|
|
parser.add_argument(
|
|
"-o", "--output",
|
|
help="key store db that will receive keys",
|
|
required=True)
|
|
parser.add_argument(
|
|
"-d", "--dry-run",
|
|
help="execute it, but never actually save/commit changes.",
|
|
action="store_true", required=False)
|
|
args = parser.parse_args()
|
|
|
|
output_db = AtomicSQL()
|
|
output_db_id = output_db.load(sqlite3.connect(args.output))
|
|
|
|
# get all keys from input db
|
|
add_count = 0
|
|
update_count = 0
|
|
existed_count = 0
|
|
|
|
if args.input == "-":
|
|
input_ = sys.stdin.read()
|
|
else:
|
|
with open(args.input, encoding="utf-8") as fd:
|
|
input_ = fd.read()
|
|
|
|
for line in input_.splitlines(keepends=False):
|
|
match = re.search(r"^(?P<kid>[0-9a-fA-F]{32}):(?P<key>[0-9a-fA-F]{32})(:(?P<title>[\w .:-]*))?$", line)
|
|
if not match:
|
|
continue
|
|
kid = match.group("kid").lower()
|
|
key = match.group("key").lower()
|
|
title = match.group("title") or None
|
|
|
|
exists = output_db.safe_execute(
|
|
output_db_id,
|
|
lambda db, cursor: cursor.execute(
|
|
f"SELECT title FROM `{args.table}` WHERE `kid`=:kid",
|
|
{"kid": kid}
|
|
)
|
|
).fetchone()
|
|
|
|
if exists:
|
|
if title and not exists[0]:
|
|
update_count += 1
|
|
print(f"Updating {args.table} {kid}: {title}")
|
|
output_db.safe_execute(
|
|
output_db_id,
|
|
lambda db, cursor: cursor.execute(
|
|
f"UPDATE `{args.table}` SET `title`=:title",
|
|
{"title": title}
|
|
)
|
|
)
|
|
else:
|
|
existed_count += 1
|
|
print(f"Key {args.table} {kid} already exists in the db with no differences, skipping...")
|
|
else:
|
|
add_count += 1
|
|
print(f"Adding {args.table} {kid} ({title}): {key}")
|
|
output_db.safe_execute(
|
|
output_db_id,
|
|
lambda db, cursor: cursor.execute(
|
|
f"INSERT INTO `{args.table}` (kid, key_, title) VALUES (:kid, :key, :title)",
|
|
{"kid": kid, "key": key, "title": title}
|
|
)
|
|
)
|
|
|
|
if args.dry_run:
|
|
print("--dry run enabled, have not commited any changes.")
|
|
else:
|
|
output_db.commit(output_db_id)
|
|
|
|
print(
|
|
"Done!\n"
|
|
f"{add_count} added, {update_count} updated in some way, {existed_count} already existed (skipped)"
|
|
)
|