Change: worker_manager to use AuthClient

1 files changed, 15 insertions, 82 deletions

diff --git a/worker_manager b/worker_manager
index fa92fa9..6fd556c 100755
--- a/worker_manager
+++ b/worker_manager
@@ -18,9 +18,7 @@
 import logging
 import time
 
-import apifw
 import cliapp
-import Crypto.PublicKey.RSA
 import urllib3
 import ick2
 
@@ -34,7 +32,6 @@ class WorkerManager(cliapp.Application):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self._token = None
-        self._token_until = None
 
     def add_settings(self):
         self.settings.string(
@@ -49,18 +46,6 @@ class WorkerManager(cliapp.Application):
             metavar='URL',
         )
 
-        self.settings.string(
-            ['token-key'],
-            'get token signing private key from FILE',
-            metavar='FILE',
-        )
-
-        self.settings.string(
-            ['token-key-pub'],
-            'this is not used',
-            metavar='NOPE',
-        )
-
         self.settings.integer(
             ['sleep'],
             'sleep for SECS seconds if there is no work currently',
@@ -105,9 +90,7 @@ class WorkerManager(cliapp.Application):
         workspace = self.settings['workspace']
         systree = self.settings['systree']
 
-        tg = TokenGenerator()
-        tg.set_key(self.settings['token-key'])
-        api = ControllerAPI(name, url, tg)
+        api = ControllerAPI(name, url)
         api.set_verify_tls(self.settings['verify-tls'])
         worker = Worker(name, api, workspace, systree)
 
@@ -128,18 +111,29 @@ class WorkerManager(cliapp.Application):
 
 class ControllerAPI:
 
-    def __init__(self, name, url, token_generator):
-        self._token_generator = token_generator
+    _scopes = ' '.join([
+        'uapi_version_get',
+        'uapi_work_id_get',
+        'uapi_work_post',
+        'uapi_workers_post',
+        'uapi_blobs_id_get',
+        'uapi_blobs_id_put',
+    ])
+
+    def __init__(self, name, url):
         self._cc = ick2.ControllerClient()
         self._cc.set_client_name(name)
         self._cc.set_controller_url(url)
+        self._ac = None
         self._blobs = None
 
     def set_verify_tls(self, verify):
         self._cc.set_verify_tls(verify)
 
     def get_token(self):
-        return self._token_generator.get_token()
+        if self._ac is None:
+            self._ac = self._cc.get_auth_client()
+        return self._ac.get_token(self._scopes)
 
     def register(self):
         self._cc.set_token(self.get_token())
@@ -171,67 +165,6 @@ class ControllerAPI:
         return self._blobs
 
 
-class TokenGenerator:
-
-    max_age = 3600  # 1 hour
-    sub = 'subject-uuid'
-    iss = 'localhost'
-    aud = 'localhost'
-    scopes = ' '.join([
-        'uapi_version_get',
-        'uapi_work_id_get',
-        'uapi_work_post',
-        'uapi_workers_post',
-        'uapi_blobs_id_get',
-        'uapi_blobs_id_put',
-    ])
-
-    def __init__(self):
-        self._token = None
-        self._token_key = None
-        self._token_until = None
-
-    def is_valid(self, now):
-        return (
-            self._token is not None and
-            (self._token_until is None or now <= self._token_until)
-        )
-
-    def set_token(self, token):
-        self._token = token
-        self._token_until = None
-        assert self.is_valid(time.time())
-
-    def set_key(self, filename):
-        key_text = self.cat(filename)
-        self._token_key = Crypto.PublicKey.RSA.importKey(key_text)
-
-    def cat(self, filename):
-        with open(filename) as f:
-            return f.read()
-
-    def get_token(self):
-        now = time.time()
-        if not self.is_valid(now):
-            self._token = self.create_token()
-            self._token_until = now + self.max_age
-        assert self.is_valid(now)
-        return self._token
-
-    def create_token(self):
-        now = time.time()
-        claims = {
-            'iss': self.iss,
-            'sub': self.sub,
-            'aud': self.aud,
-            'exp': now + self.max_age,
-            'scope': self.scopes,
-        }
-
-        token = apifw.create_token(claims, self._token_key)
-        return token.decode('ascii')
-
-
 class Worker:
 
     def __init__(self, name, api, workspace, systree):