added basic sched auth

daskos · Apr 13, 2017 · 34e2467 · 34e2467
1 parent 9d34b40
commit 34e2467
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 4 deletions.
diff --git a/mentos/connection.py b/mentos/connection.py
@@ -8,6 +8,7 @@
                                OutBoundError)
 from mentos.utils import decode, encode, log_errors
 from six import raise_from
+from binascii import b2a_base64
 from six.moves.urllib.parse import urlparse
 from tornado import concurrent, gen
 from tornado.httpclient import AsyncHTTPClient, HTTPError, HTTPRequest
@@ -26,7 +27,7 @@ class Connection(object):
         'Connection': 'close'
     }
 
-    def __init__(self, endpoint, api_path, event_handler):
+    def __init__(self, endpoint, api_path, event_handler, principal=None, secret=None):
 
         self.endpoint = endpoint
         self.api_path = api_path
@@ -44,6 +45,13 @@ def __init__(self, endpoint, api_path, event_handler):
         self.connection_successful = False
         self._headers = HTTPHeaders()
 
+        if principal is not None and secret is not None:
+            self.headers['Authorization'] = self._basic_credential = 'Basic %s' % (
+                b2a_base64(
+                    ('%s:%s' % (principal, secret)).encode('ascii')
+                ).decode('ascii').strip()
+            )
+
     def _parse_subscription_headers(self, response):
         try:
             if "200 OK" in response:

diff --git a/mentos/scheduler.py b/mentos/scheduler.py
@@ -19,7 +19,7 @@ class SchedulerDriver(object):
 
     def __init__(self, scheduler, name, user=getpass.getuser(),
                  master=os.getenv('MESOS_MASTER', 'zk://localhost:2181'),
-                 failover_timeout=100, capabilities=None,
+                 failover_timeout=100, capabilities=None, principal=None, secret=None,
                  implicit_acknowledgements=True, handlers={}, loop=None):
         self.loop = loop or IOLoop()
         self.master = master
@@ -52,6 +52,8 @@ def __init__(self, scheduler, name, user=getpass.getuser(),
 
         self.subscription = Subscription(self.framework, self.master,
                                          '/api/v1/scheduler', self.handlers,
+                                         principal=principal,
+                                         secret=secret,
                                          timeout=failover_timeout,
                                          loop=self.loop)
 

diff --git a/mentos/subscription.py b/mentos/subscription.py
@@ -53,7 +53,7 @@ class Message(object):
 
 class Subscription(object):
 
-    def __init__(self, framework, uri, api_path, event_handlers=None,
+    def __init__(self, framework, uri, api_path, event_handlers=None, principal=None, secret=None,
                  timeout=75, retry_policy=RetryPolicy.forever(), loop=None):
         self.loop = loop or IOLoop.current()
 
@@ -79,6 +79,9 @@ def __init__(self, framework, uri, api_path, event_handlers=None,
         self.updates = None
         self.timeout = timeout
 
+        self.principal = principal
+        self.secret = secret
+
     @gen.coroutine
     def ensure_safe(self, safe_states=[States.SUBSCRIBED, States.SUBSCRIBING]):
         if self.state in safe_states:
@@ -180,7 +183,8 @@ def subscribe(self):
 
     @gen.coroutine
     def make_connection(self, endpoint, api_path):
-        conn = Connection(endpoint, api_path, self._event_handler)
+        conn = Connection(endpoint, api_path, self._event_handler,
+                          principal=self.principal, secret=self.secret)
         try:
             yield conn.ping()
         except MasterRedirect as ex:  # pragma: no cover