Merge pull request #52 from windreamer/dev

support http basic authentication

pymesos/scheduler.py

@@ -2,6 +2,7 @@
 import logging
 from addict import Dict
 from six.moves.http_client import HTTPConnection
+from binascii import b2a_base64
 from .process import Process
 from .interface import SchedulerDriver
 
@@ -12,7 +13,8 @@ class MesosSchedulerDriver(Process, SchedulerDriver):
     _timeout = 10
 
     def __init__(self, sched, framework, master_uri,
-                 use_addict=False, implicit_acknowledgements=True):
+                 use_addict=False, implicit_acknowledgements=True,
+                 principal=None, secret=None):
         super(MesosSchedulerDriver, self).__init__()
         self.sched = sched
         self.master_uri = master_uri
@@ -23,6 +25,14 @@ def __init__(self, sched, framework, master_uri,
         self._failover = False
         self._dict_cls = Dict if use_addict else dict
         self.implicit_acknowledgements = implicit_acknowledgements
+        if principal is not None and secret is not None:
+            self._basic_credential = 'Basic %s' % (
+                b2a_base64(
+                    ('%s:%s' % (principal, secret)).encode('ascii')
+                ).decode('ascii').strip()
+            )
+        else:
+            self._basic_credential = None
 
     @property
     def framework(self):
@@ -141,6 +151,9 @@ def _send(self, body, path='/api/v1/scheduler', method='POST', headers={}):
             if stream_id:
                 headers['Mesos-Stream-Id'] = stream_id
 
+            if self._basic_credential:
+                headers['Authorization'] = self._basic_credential
+
             try:
                 conn.request(method, path, body=data, headers=headers)
                 resp = conn.getresponse()
@@ -353,11 +366,17 @@ def gen_request(self):
             request['framework_id'] = self._framework['id']
 
         data = json.dumps(request)
+        _authorization = ''
+        if self._basic_credential is not None:
+            _authorization = 'Authorization: %s\r\n' % (
+                self._basic_credential,
+            )
+
         request = ('POST /api/v1/scheduler HTTP/1.1\r\nHost: %s\r\n'
                    'Content-Type: application/json\r\n'
-                   'Accept: application/json\r\n'
+                   'Accept: application/json\r\n%s'
                    'Connection: close\r\nContent-Length: %s\r\n\r\n%s') % (
-                       self.master, len(data), data
+                       self.master, _authorization, len(data), data
         )
         return request.encode('utf-8')
 

pymesos/subprocess/scheduler.py

@@ -23,7 +23,16 @@ def __init__(self):
         self.framework = self._init_framework()
         self.executor = None
         self.master = str(CONFIG.get('master', os.environ['MESOS_MASTER']))
-        self.driver = MesosSchedulerDriver(self, self.framework, self.master)
+        self.principal = str(
+            CONFIG.get('principal', os.environ.get('DEFAULT_PRINCIPAL'))
+        )
+        self.secret = str(
+            CONFIG.get('secret', os.environ.get('DEFAULT_SECRET'))
+        )
+        self.driver = MesosSchedulerDriver(
+            self, self.framework, self.master, principal=self.principal,
+            secret=self.secret
+        )
         self.procs_pending = {}
         self.procs_launched = {}
         self.agent_to_proc = {}