#! /usr/bin/env python3 # # Copyright (C) 2019 Garmin Ltd. # # SPDX-License-Identifier: GPL-2.0-only # import argparse import hashlib import logging import os import pprint import sys import threading import time import warnings import netrc import json warnings.simplefilter("default") try: import tqdm ProgressBar = tqdm.tqdm except ImportError: class ProgressBar(object): def __init__(self, *args, **kwargs): pass def __enter__(self): return self def __exit__(self, *args, **kwargs): pass def update(self): pass sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(__file__)), 'lib')) import hashserv import bb.asyncrpc DEFAULT_ADDRESS = 'unix://./hashserve.sock' METHOD = 'stress.test.method' def print_user(u): print(f"Username: {u['username']}") if "permissions" in u: print("Permissions: " + " ".join(u["permissions"])) if "token" in u: print(f"Token: {u['token']}") def main(): def handle_get(args, client): result = client.get_taskhash(args.method, args.taskhash, all_properties=True) if not result: return 0 print(json.dumps(result, sort_keys=True, indent=4)) return 0 def handle_get_outhash(args, client): result = client.get_outhash(args.method, args.outhash, args.taskhash) if not result: return 0 print(json.dumps(result, sort_keys=True, indent=4)) return 0 def handle_stats(args, client): if args.reset: s = client.reset_stats() else: s = client.get_stats() print(json.dumps(s, sort_keys=True, indent=4)) return 0 def handle_stress(args, client): def thread_main(pbar, lock): nonlocal found_hashes nonlocal missed_hashes nonlocal max_time with hashserv.create_client(args.address) as client: for i in range(args.requests): taskhash = hashlib.sha256() taskhash.update(args.taskhash_seed.encode('utf-8')) taskhash.update(str(i).encode('utf-8')) start_time = time.perf_counter() l = client.get_unihash(METHOD, taskhash.hexdigest()) elapsed = time.perf_counter() - start_time with lock: if l: found_hashes += 1 else: missed_hashes += 1 max_time = max(elapsed, max_time) pbar.update() max_time = 0 found_hashes = 0 missed_hashes = 0 lock = threading.Lock() total_requests = args.clients * args.requests start_time = time.perf_counter() with ProgressBar(total=total_requests) as pbar: threads = [threading.Thread(target=thread_main, args=(pbar, lock), daemon=False) for _ in range(args.clients)] for t in threads: t.start() for t in threads: t.join() elapsed = time.perf_counter() - start_time with lock: print("%d requests in %.1fs. %.1f requests per second" % (total_requests, elapsed, total_requests / elapsed)) print("Average request time %.8fs" % (elapsed / total_requests)) print("Max request time was %.8fs" % max_time) print("Found %d hashes, missed %d" % (found_hashes, missed_hashes)) if args.report: with ProgressBar(total=args.requests) as pbar: for i in range(args.requests): taskhash = hashlib.sha256() taskhash.update(args.taskhash_seed.encode('utf-8')) taskhash.update(str(i).encode('utf-8')) outhash = hashlib.sha256() outhash.update(args.outhash_seed.encode('utf-8')) outhash.update(str(i).encode('utf-8')) client.report_unihash(taskhash.hexdigest(), METHOD, outhash.hexdigest(), taskhash.hexdigest()) with lock: pbar.update() def handle_remove(args, client): where = {k: v for k, v in args.where} if where: result = client.remove(where) print("Removed %d row(s)" % (result["count"])) else: print("No query specified") def handle_clean_unused(args, client): result = client.clean_unused(args.max_age) print("Removed %d rows" % (result["count"])) return 0 def handle_refresh_token(args, client): r = client.refresh_token(args.username) print_user(r) def handle_set_user_permissions(args, client): r = client.set_user_perms(args.username, args.permissions) print_user(r) def handle_get_user(args, client): r = client.get_user(args.username) print_user(r) def handle_get_all_users(args, client): users = client.get_all_users() print("{username:20}| {permissions}".format(username="Username", permissions="Permissions")) print(("-" * 20) + "+" + ("-" * 20)) for u in users: print("{username:20}| {permissions}".format(username=u["username"], permissions=" ".join(u["permissions"]))) def handle_new_user(args, client): r = client.new_user(args.username, args.permissions) print_user(r) def handle_delete_user(args, client): r = client.delete_user(args.username) print_user(r) def handle_get_db_usage(args, client): usage = client.get_db_usage() print(usage) tables = sorted(usage.keys()) print("{name:20}| {rows:20}".format(name="Table name", rows="Rows")) print(("-" * 20) + "+" + ("-" * 20)) for t in tables: print("{name:20}| {rows:<20}".format(name=t, rows=usage[t]["rows"])) print() total_rows = sum(t["rows"] for t in usage.values()) print(f"Total rows: {total_rows}") def handle_get_db_query_columns(args, client): columns = client.get_db_query_columns() print("\n".join(sorted(columns))) def handle_gc_status(args, client): result = client.gc_status() if not result["mark"]: print("No Garbage collection in progress") return 0 print("Current Mark: %s" % result["mark"]) print("Total hashes to keep: %d" % result["keep"]) print("Total hashes to remove: %s" % result["remove"]) return 0 def handle_gc_mark(args, client): where = {k: v for k, v in args.where} result = client.gc_mark(args.mark, where) print("New hashes marked: %d" % result["count"]) return 0 def handle_gc_sweep(args, client): result = client.gc_sweep(args.mark) print("Removed %d rows" % result["count"]) return 0 def handle_unihash_exists(args, client): result = client.unihash_exists(args.unihash) if args.quiet: return 0 if result else 1 print("true" if result else "false") return 0 parser = argparse.ArgumentParser(description='Hash Equivalence Client') parser.add_argument('--address', default=DEFAULT_ADDRESS, help='Server address (default "%(default)s")') parser.add_argument('--log', default='WARNING', help='Set logging level') parser.add_argument('--login', '-l', metavar="USERNAME", help="Authenticate as USERNAME") parser.add_argument('--password', '-p', metavar="TOKEN", help="Authenticate using token TOKEN") parser.add_argument('--become', '-b', metavar="USERNAME", help="Impersonate user USERNAME (if allowed) when performing actions") parser.add_argument('--no-netrc', '-n', action="store_false", dest="netrc", help="Do not use .netrc") subparsers = parser.add_subparsers() get_parser = subparsers.add_parser('get', help="Get the unihash for a taskhash") get_parser.add_argument("method", help="Method to query") get_parser.add_argument("taskhash", help="Task hash to query") get_parser.set_defaults(func=handle_get) get_outhash_parser = subparsers.add_parser('get-outhash', help="Get output hash information") get_outhash_parser.add_argument("method", help="Method to query") get_outhash_parser.add_argument("outhash", help="Output hash to query") get_outhash_parser.add_argument("taskhash", help="Task hash to query") get_outhash_parser.set_defaults(func=handle_get_outhash) stats_parser = subparsers.add_parser('stats', help='Show server stats') stats_parser.add_argument('--reset', action='store_true', help='Reset server stats') stats_parser.set_defaults(func=handle_stats) stress_parser = subparsers.add_parser('stress', help='Run stress test') stress_parser.add_argument('--clients', type=int, default=10, help='Number of simultaneous clients') stress_parser.add_argument('--requests', type=int, default=1000, help='Number of requests each client will perform') stress_parser.add_argument('--report', action='store_true', help='Report new hashes') stress_parser.add_argument('--taskhash-seed', default='', help='Include string in taskhash') stress_parser.add_argument('--outhash-seed', default='', help='Include string in outhash') stress_parser.set_defaults(func=handle_stress) remove_parser = subparsers.add_parser('remove', help="Remove hash entries") remove_parser.add_argument("--where", "-w", metavar="KEY VALUE", nargs=2, action="append", default=[], help="Remove entries from table where KEY == VALUE") remove_parser.set_defaults(func=handle_remove) clean_unused_parser = subparsers.add_parser('clean-unused', help="Remove unused database entries") clean_unused_parser.add_argument("max_age", metavar="SECONDS", type=int, help="Remove unused entries older than SECONDS old") clean_unused_parser.set_defaults(func=handle_clean_unused) refresh_token_parser = subparsers.add_parser('refresh-token', help="Refresh auth token") refresh_token_parser.add_argument("--username", "-u", help="Refresh the token for another user (if authorized)") refresh_token_parser.set_defaults(func=handle_refresh_token) set_user_perms_parser = subparsers.add_parser('set-user-perms', help="Set new permissions for user") set_user_perms_parser.add_argument("--username", "-u", help="Username", required=True) set_user_perms_parser.add_argument("permissions", metavar="PERM", nargs="*", default=[], help="New permissions") set_user_perms_parser.set_defaults(func=handle_set_user_permissions) get_user_parser = subparsers.add_parser('get-user', help="Get user") get_user_parser.add_argument("--username", "-u", help="Username") get_user_parser.set_defaults(func=handle_get_user) get_all_users_parser = subparsers.add_parser('get-all-users', help="List all users") get_all_users_parser.set_defaults(func=handle_get_all_users) new_user_parser = subparsers.add_parser('new-user', help="Create new user") new_user_parser.add_argument("--username", "-u", help="Username", required=True) new_user_parser.add_argument("permissions", metavar="PERM", nargs="*", default=[], help="New permissions") new_user_parser.set_defaults(func=handle_new_user) delete_user_parser = subparsers.add_parser('delete-user', help="Delete user") delete_user_parser.add_argument("--username", "-u", help="Username", required=True) delete_user_parser.set_defaults(func=handle_delete_user) db_usage_parser = subparsers.add_parser('get-db-usage', help="Database Usage") db_usage_parser.set_defaults(func=handle_get_db_usage) db_query_columns_parser = subparsers.add_parser('get-db-query-columns', help="Show columns that can be used in database queries") db_query_columns_parser.set_defaults(func=handle_get_db_query_columns) gc_status_parser = subparsers.add_parser("gc-status", help="Show garbage collection status") gc_status_parser.set_defaults(func=handle_gc_status) gc_mark_parser = subparsers.add_parser('gc-mark', help="Mark hashes to be kept for garbage collection") gc_mark_parser.add_argument("mark", help="Mark for this garbage collection operation") gc_mark_parser.add_argument("--where", "-w", metavar="KEY VALUE", nargs=2, action="append", default=[], help="Keep entries in table where KEY == VALUE") gc_mark_parser.set_defaults(func=handle_gc_mark) gc_sweep_parser = subparsers.add_parser('gc-sweep', help="Perform garbage collection and delete any entries that are not marked") gc_sweep_parser.add_argument("mark", help="Mark for this garbage collection operation") gc_sweep_parser.set_defaults(func=handle_gc_sweep) unihash_exists_parser = subparsers.add_parser('unihash-exists', help="Check if a unihash is known to the server") unihash_exists_parser.add_argument("--quiet", action="store_true", help="Don't print status. Instead, exit with 0 if unihash exists and 1 if it does not") unihash_exists_parser.add_argument("unihash", help="Unihash to check") unihash_exists_parser.set_defaults(func=handle_unihash_exists) args = parser.parse_args() logger = logging.getLogger('hashserv') level = getattr(logging, args.log.upper(), None) if not isinstance(level, int): raise ValueError('Invalid log level: %s' % args.log) logger.setLevel(level) console = logging.StreamHandler() console.setLevel(level) logger.addHandler(console) login = args.login password = args.password if login is None and args.netrc: try: n = netrc.netrc() auth = n.authenticators(args.address) if auth is not None: login, _, password = auth except FileNotFoundError: pass except netrc.NetrcParseError as e: sys.stderr.write(f"Error parsing {e.filename}:{e.lineno}: {e.msg}\n") func = getattr(args, 'func', None) if func: try: with hashserv.create_client(args.address, login, password) as client: if args.become: client.become_user(args.become) return func(args, client) except bb.asyncrpc.InvokeError as e: print(f"ERROR: {e}") return 1 return 0 if __name__ == '__main__': try: ret = main() except Exception: ret = 1 import traceback traceback.print_exc() sys.exit(ret)