diff options
| -rwxr-xr-x | Run/Que/client.py | 26 |
1 files changed, 21 insertions, 5 deletions
diff --git a/Run/Que/client.py b/Run/Que/client.py index 4cf5d05..3d9291d 100755 --- a/Run/Que/client.py +++ b/Run/Que/client.py @@ -32,7 +32,7 @@ def auth(args): def send(args): "Send a message to the que." key = auth(args) - data = args.infile.read().encode("utf-8").strip() + data = args.infile req = request.Request(f"{args.host}/{args.target}") req.add_header("User-AgenT", "Que/Client") if key: @@ -47,10 +47,9 @@ def send(args): def recv(args): "Receive a message from the que." - key = auth(args) def _recv(_req): - msg = _req.readline().decode("utf-8").strip() + msg = autodecode(_req.read()) print(msg) if args.then: subprocess.run( @@ -60,9 +59,9 @@ def recv(args): params = urllib.parse.urlencode({"poll": args.poll}) req = request.Request(f"{args.host}/{args.target}?{params}") req.add_header("User-Agent", "Que/Client") + key = auth(args) if key: req.add_header("Authorization", key) - with request.urlopen(req) as _req: if args.poll: while not time.sleep(1): @@ -71,6 +70,23 @@ def recv(args): _recv(_req) +def autodecode(b): + """Attempt to decode bytes `b` into common codecs, preferably utf-8. If + no decoding is available, just return the raw bytes. + + For all available codecs, see: + <https://docs.python.org/3/library/codecs.html#standard-encodings> + + """ + codecs = ["utf-8", "ascii"] + for codec in codecs: + try: + return b.decode(codec) + except UnicodeDecodeError: + pass + return b + + def get_args(): cli = argparse.ArgumentParser(description=__doc__) cli.add_argument( @@ -106,7 +122,7 @@ def get_args(): cli.add_argument( "infile", nargs="?", - type=argparse.FileType("r"), + type=argparse.FileType("rb"), help="data to put on the que. Use '-' for stdin, otherwise should be a readable file", ) return cli.parse_args() |