X-Git-Url: http://gitweb.pimeys.fr/?p=bots%2Fbasile.git;a=blobdiff_plain;f=nk.py;h=0aafbc16aff64606b07b40e5d5ccf0e6a6ea6e69;hp=a9e8a5c9d923b11d27cbb596cdd9e9e015bf7a2d;hb=HEAD;hpb=b88d22150a746bb953ca9500e95a391c20f377a6 diff --git a/nk.py b/nk.py index a9e8a5c..0aafbc1 100644 --- a/nk.py +++ b/nk.py @@ -34,19 +34,25 @@ class NKUnknownError(NKError): class NKDeadServer(NKError): pass -def full_read(sock): - """Lit sur la socket jusqu'à ce que l'output soit déJSON-izable""" - output = "" - i = 0 - while True: - output += sock.read() - i += 1 - try: - return json.loads(output) - except: - pass - if i == 10 and output == "": - raise NKDeadServer +def full_read(socket): + # On récupère d'abord la taille du message + length_str = '' + char = socket.recv(1) + while char != '\n': + length_str += char + char = socket.recv(1) + total = int(length_str) + # On utilise une memoryview pour recevoir les données chunk par chunk efficacement + view = memoryview(bytearray(total)) + next_offset = 0 + while total - next_offset > 0: + recv_size = socket.recv_into(view[next_offset:], total - next_offset) + next_offset += recv_size + try: + msg = json.loads(view.tobytes()) + except (TypeError, ValueError) as e: + raise NKNotJson("L'objet reçu n'est pas un JSON") + return msg def connect(): sock = socket.socket() @@ -86,25 +92,29 @@ def login(username, password, typ="bdd"): raise NKRefused(str(exc)) return out, sock -def get_solde(sock, idbde, serv, canal): - """Récupère le (success, solde, pseudo) de l'utilisateur NK n°``idbde``""" +def get_infos(sock, idbde, serv, canal): + """Récupère les données de l'utilisateur NK n°``idbde``""" try: sock.write(json.dumps(["compte", idbde])) ret = full_read(sock) retcode = ret["retcode"] if retcode == 0: - solde = ret["msg"]["solde"] - pseudo = ret["msg"]["pseudo"] - return (True, solde, pseudo) + return ret["msg"] else: serv.privmsg(canal, ret["errmsg"].encode("utf-8")) - return (False, None, None) except Exception as exc: trace = traceback.format_exc() msg = "failed\n%s" % trace for l in msg.split("\n"): serv.privmsg(canal, l) #log(self.serveur, "priv", auteur, " ".join(message) + "[failed]") + +def get_solde(sock, idbde, serv, canal): + """Récupère le (success, solde, pseudo) de l'utilisateur NK n°``idbde``""" + infos = get_infos(sock, idbde, serv, canal) + if infos: + return (True, infos["solde"], infos["pseudo"]) + else: return (False, None, None) def consomme(sock, idbde, conso, serv, canal):