Add auth header to que client

1 files changed, 19 insertions, 0 deletions

diff --git a/Run/Que/client.py b/Run/Que/client.py
index 43517ff..ed128b8 100755
--- a/Run/Que/client.py
+++ b/Run/Que/client.py
@@ -4,6 +4,8 @@ simple client for que.run
 """
 
 import argparse
+import configparser
+import os
 import subprocess
 import sys
 import time
@@ -13,11 +15,24 @@ import urllib.request as request
 MAX_TIMEOUT = 99999999  # basically never timeout
 
 
+def auth():
+    conf_file = os.path.expanduser("~/.config/que.conf")
+    if not os.path.exists(conf_file):
+        sys.exit("you need a ~/.config/que.conf")
+    cfg = configparser.ConfigParser()
+    cfg.read(conf_file)
+    return cfg
+
+
 def send(args):
     "Send a message to the que."
+    ns = args.target.split("/")[0]
+    key = auth()[ns]["key"]
     data = args.infile.read().encode("utf-8").strip()
     req = request.Request(f"{args.host}/{args.target}")
     req.add_header("User-AgenT", "Que/Client")
+    if key:
+        req.add_header("Authorization", key)
     if args.serve:
         while not time.sleep(1):
             with request.urlopen(req, data=data, timeout=MAX_TIMEOUT) as req:
@@ -29,6 +44,8 @@ def send(args):
 
 def recv(args):
     "Receive a message from the que."
+    ns = args.target.split("/")[0]
+    key = auth()[ns]["key"]
 
     def _recv(_req):
         msg = _req.readline().decode("utf-8").strip()
@@ -41,6 +58,8 @@ 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")
+    if key:
+        req.add_header("Authorization", key)
 
     with request.urlopen(req) as _req:
         if args.poll: