#!/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 `:` 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)" )