# Copyright (C) 2019 Garmin Ltd. # # SPDX-License-Identifier: GPL-2.0-only # import logging import socket import bb.asyncrpc import json from . import create_async_client logger = logging.getLogger("hashserv.client") class AsyncClient(bb.asyncrpc.AsyncClient): MODE_NORMAL = 0 MODE_GET_STREAM = 1 def __init__(self, username=None, password=None): super().__init__("OEHASHEQUIV", "1.1", logger) self.mode = self.MODE_NORMAL self.username = username self.password = password self.saved_become_user = None async def setup_connection(self): await super().setup_connection() cur_mode = self.mode self.mode = self.MODE_NORMAL await self._set_mode(cur_mode) if self.username: # Save off become user temporarily because auth() resets it become = self.saved_become_user await self.auth(self.username, self.password) if become: await self.become_user(become) async def send_stream(self, msg): async def proc(): await self.socket.send(msg) return await self.socket.recv() return await self._send_wrapper(proc) async def _set_mode(self, new_mode): async def stream_to_normal(): await self.socket.send("END") return await self.socket.recv() if new_mode == self.MODE_NORMAL and self.mode == self.MODE_GET_STREAM: r = await self._send_wrapper(stream_to_normal) if r != "ok": self.check_invoke_error(r) raise ConnectionError("Unable to transition to normal mode: Bad response from server %r" % r) elif new_mode == self.MODE_GET_STREAM and self.mode == self.MODE_NORMAL: r = await self.invoke({"get-stream": None}) if r != "ok": raise ConnectionError("Unable to transition to stream mode: Bad response from server %r" % r) elif new_mode != self.mode: raise Exception( "Undefined mode transition %r -> %r" % (self.mode, new_mode) ) self.mode = new_mode async def get_unihash(self, method, taskhash): await self._set_mode(self.MODE_GET_STREAM) r = await self.send_stream("%s %s" % (method, taskhash)) if not r: return None return r async def report_unihash(self, taskhash, method, outhash, unihash, extra={}): await self._set_mode(self.MODE_NORMAL) m = extra.copy() m["taskhash"] = taskhash m["method"] = method m["outhash"] = outhash m["unihash"] = unihash return await self.invoke({"report": m}) async def report_unihash_equiv(self, taskhash, method, unihash, extra={}): await self._set_mode(self.MODE_NORMAL) m = extra.copy() m["taskhash"] = taskhash m["method"] = method m["unihash"] = unihash return await self.invoke({"report-equiv": m}) async def get_taskhash(self, method, taskhash, all_properties=False): await self._set_mode(self.MODE_NORMAL) return await self.invoke( {"get": {"taskhash": taskhash, "method": method, "all": all_properties}} ) async def get_outhash(self, method, outhash, taskhash, with_unihash=True): await self._set_mode(self.MODE_NORMAL) return await self.invoke( { "get-outhash": { "outhash": outhash, "taskhash": taskhash, "method": method, "with_unihash": with_unihash, } } ) async def get_stats(self): await self._set_mode(self.MODE_NORMAL) return await self.invoke({"get-stats": None}) async def reset_stats(self): await self._set_mode(self.MODE_NORMAL) return await self.invoke({"reset-stats": None}) async def backfill_wait(self): await self._set_mode(self.MODE_NORMAL) return (await self.invoke({"backfill-wait": None}))["tasks"] async def remove(self, where): await self._set_mode(self.MODE_NORMAL) return await self.invoke({"remove": {"where": where}}) async def clean_unused(self, max_age): await self._set_mode(self.MODE_NORMAL) return await self.invoke({"clean-unused": {"max_age_seconds": max_age}}) async def auth(self, username, token): await self._set_mode(self.MODE_NORMAL) result = await self.invoke({"auth": {"username": username, "token": token}}) self.username = username self.password = token self.saved_become_user = None return result async def refresh_token(self, username=None): await self._set_mode(self.MODE_NORMAL) m = {} if username: m["username"] = username result = await self.invoke({"refresh-token": m}) if ( self.username and not self.saved_become_user and result["username"] == self.username ): self.password = result["token"] return result async def set_user_perms(self, username, permissions): await self._set_mode(self.MODE_NORMAL) return await self.invoke( {"set-user-perms": {"username": username, "permissions": permissions}} ) async def get_user(self, username=None): await self._set_mode(self.MODE_NORMAL) m = {} if username: m["username"] = username return await self.invoke({"get-user": m}) async def get_all_users(self): await self._set_mode(self.MODE_NORMAL) return (await self.invoke({"get-all-users": {}}))["users"] async def new_user(self, username, permissions): await self._set_mode(self.MODE_NORMAL) return await self.invoke( {"new-user": {"username": username, "permissions": permissions}} ) async def delete_user(self, username): await self._set_mode(self.MODE_NORMAL) return await self.invoke({"delete-user": {"username": username}}) async def become_user(self, username): await self._set_mode(self.MODE_NORMAL) result = await self.invoke({"become-user": {"username": username}}) if username == self.username: self.saved_become_user = None else: self.saved_become_user = username return result async def get_db_usage(self): await self._set_mode(self.MODE_NORMAL) return (await self.invoke({"get-db-usage": {}}))["usage"] async def get_db_query_columns(self): await self._set_mode(self.MODE_NORMAL) return (await self.invoke({"get-db-query-columns": {}}))["columns"] async def gc_status(self): await self._set_mode(self.MODE_NORMAL) return await self.invoke({"gc-status": {}}) async def gc_mark(self, mark, where): """ Starts a new garbage collection operation identified by "mark". If garbage collection is already in progress with "mark", the collection is continued. All unihash entries that match the "where" clause are marked to be kept. In addition, any new entries added to the database after this command will be automatically marked with "mark" """ await self._set_mode(self.MODE_NORMAL) return await self.invoke({"gc-mark": {"mark": mark, "where": where}}) async def gc_sweep(self, mark): """ Finishes garbage collection for "mark". All unihash entries that have not been marked will be deleted. It is recommended to clean unused outhash entries after running this to cleanup any dangling outhashes """ await self._set_mode(self.MODE_NORMAL) return await self.invoke({"gc-sweep": {"mark": mark}}) class Client(bb.asyncrpc.Client): def __init__(self, username=None, password=None): self.username = username self.password = password super().__init__() self._add_methods( "connect_tcp", "connect_websocket", "get_unihash", "report_unihash", "report_unihash_equiv", "get_taskhash", "get_outhash", "get_stats", "reset_stats", "backfill_wait", "remove", "clean_unused", "auth", "refresh_token", "set_user_perms", "get_user", "get_all_users", "new_user", "delete_user", "become_user", "get_db_usage", "get_db_query_columns", "gc_status", "gc_mark", "gc_sweep", ) def _get_async_client(self): return AsyncClient(self.username, self.password)