Add: create token in worker-manager

1 files changed, 52 insertions, 1 deletions

diff --git a/worker_manager b/worker_manager
index f5492cb..1b3bbbb 100755
--- a/worker_manager
+++ b/worker_manager
@@ -20,7 +20,9 @@ import logging
 import sys
 import time
 
+import apifw
 import cliapp
+import Crypto.PublicKey.RSA
 import requests
 import urllib3
 
@@ -57,6 +59,18 @@ class WorkerManager(cliapp.Application):
             metavar='TOKEN',
         )
 
+        self.settings.string(
+            ['token-key'],
+            'get token signing private key from FILE',
+            metavar='FILE',
+        )
+
+        self.settings.string(
+            ['token-key-pub'],
+            'get token signing public key from FILE',
+            metavar='FILE',
+        )
+
         self.settings.integer(
             ['sleep'],
             'sleep for SECS seconds if there is no work currently',
@@ -97,11 +111,48 @@ class WorkerManager(cliapp.Application):
         return work
 
     def get_auth_headers(self):
-        token = self.settings['token']
+        token = self.get_token()
         return {
             'Authorization': 'Bearer {}'.format(token),
         }
 
+    def get_token(self):
+        token = self.settings['token']
+        token_key = self.settings['token-key']
+        token_key_pub = self.settings['token-key-pub']
+
+        if token:
+            return token
+        elif token_key and token_key_pub:
+            with open(token_key) as f1, open(token_key_pub) as f2:
+                key_text = f1.read()
+                pub_text = f2.read()
+                return self.create_token(key_text, pub_text)
+        sys.exit('No token and no way to create')
+
+    def create_token(self, key_text, pub_text):
+        iss = 'localhost'
+        aud = 'localhost'
+        scopes_list = [
+            'uapi_work_id_get',
+            'uapi_work_id_post',
+        ]
+
+        key = Crypto.PublicKey.RSA.importKey(key_text)
+        scopes = ' '.join(scopes_list)
+
+        now = time.time()
+        claims = {
+            'iss': iss,
+            'sub': 'subject-uuid',
+            'aud': aud,
+            'exp': now + 86400,
+            'scope': scopes,
+        }
+
+        token = apifw.create_token(claims, key)
+        return token.decode('ascii')
+
     def report_pretend_work(self, url, name, work):
         self.show_msg('Pretending to work: {!r}'.format(work))
         snippet_url = '{}/work/{}'.format(url, name)