def connection_watcher(self, h, type, state, path):
self.handle = h
self.conn_cv.acquire()
self.connected = True
zookeeper.add_auth(self.handle, "digest", "super:secret", None)
self.conn_cv.notifyAll()
self.conn_cv.release()
def maybe_authenticate():
if self._authenticated.is_set() or not self._credentials:
activate()
return
try:
scheme, credentials = self._credentials
zookeeper.add_auth(self._zh, scheme, credentials, on_authentication)
except zookeeper.ZooKeeperException as e:
self._logger('Failed to authenticate: %s' % e)
def add_auth(self, scheme, id, callback=None):
"""
添加认证
"""
zookeeper.add_auth(self.handle, scheme, id, callback)
try:
self.get("/")
return 1
except zookeeper.NoAuthException:
return 0
def maybe_authenticate():
if self._authenticated.is_set() or not self._credentials:
activate()
return
try:
scheme, credentials = self._credentials
zookeeper.add_auth(self._zh, scheme, credentials,
on_authentication)
except zookeeper.ZooKeeperException as e:
self._logger('Failed to authenticate: %s' % e)
def add_auth(self, scheme, id, callback=None):
"""
添加认证
"""
zookeeper.add_auth(self.handle, scheme, id, callback)
try:
self.get("/")
return 1
except zookeeper.NoAuthException:
return 0
def start(self):
conn_cv.acquire()
self.zh = zookeeper.init(self.zkservers, self.conn_callback, session_time)
conn_cv.wait()
conn_cv.release()
if self.credential!=None:
print "credential: "+self.credential
conn_cv.acquire()
zookeeper.add_auth(self.zh, "digest", self.credential, self.auth_callback)
conn_cv.wait()
conn_cv.release()
logger.info("ZooKeeper connection established.")
def add_auth_async(self, scheme, credential):
async_result = self._new_async_result()
def callback(handle, code):
if code != zookeeper.OK:
exc = err_to_exception(code)
async_result.set_exception(exc)
else:
async_result.set(None)
zookeeper.add_auth(self._handle, scheme, credential, callback)
return async_result
def start(self):
conn_cv.acquire()
self.zh = zookeeper.init(self.zkservers, self.conn_callback, session_time)
conn_cv.wait()
conn_cv.release()
if self.credential!=None:
print "credential: "+self.credential
conn_cv.acquire()
zookeeper.add_auth(self.zh, "digest", self.credential, self.auth_callback)
conn_cv.wait()
conn_cv.release()
logger.info("ZooKeeper connection established.")
def add_auth_async(self, scheme, credential):
"""Asynchronously send credentials to server
@param scheme: authentication scheme (default supported: "digest")
@param credential: the credential -- value depends on scheme
@return: AsyncResult object set on completion
@rtype AsyncResult
"""
async_result = self._sync.async_result()
callback = partial(_generic_callback, async_result)
zookeeper.add_auth(self._handle, scheme, credential, callback)
return async_result
def add_auth_async(self, scheme, credential):
"""Asynchronously send credentials to server
@param scheme: authentication scheme (default supported: "digest")
@param credential: the credential -- value depends on scheme
@return: AsyncResult object set on completion
@rtype AsyncResult
"""
async_result = self._sync.async_result()
callback = partial(_generic_callback, async_result)
zookeeper.add_auth(self._handle, scheme, credential, callback)
return async_result
def safe_connect(self):
self._conn_cv.acquire()
self._handler = zookeeper.init(
ZKConf.ZK_HOST, self.watcher_conn,
ZKConf.ZOO_SESSION_TIMEOUT) # milliseconds
self._conn_cv.wait(8.0)
self._conn_cv.release()
try_times = 10
while not self._connected and try_times > 0:
self.connect()
try_times -= 1
if not self._connected:
raise ZooKeeperError("zookeeper [ %s ] safe_connect() failed" %
str(ZKConf.ZK_HOST))
try:
user_pass = "******" % (ZKConf.ZOO_USERNAME, ZKConf.ZOO_PASSWORD)
utils.log(utils.cur(), self._name, "--add_auth--")
zookeeper.add_auth(self._handler, "digest", user_pass, None)
if self.register_node_cb:
try:
func = getattr(self, self.register_node_cb)
if func and callable(func):
utils.log(utils.cur(), self.register_node_cb,
"%s callback %s" % ("*" * 10, "*" * 10))
def _do_register_node():
func(*self._args, **self._kwargs)
_do_register_node()
"""
check whether temporary nodes are still existed
when session timeout elapsed
because there are old temporary nodes when register new
temporary nodes
"""
threading.Timer(int(ZKConf.ZOO_SESSION_TIMEOUT / 1000),
_do_register_node).start()
except AttributeError, err:
utils.err(utils.cur(), traceback.format_exc())
utils.log(utils.cur(), self._first, self.register_watch_cb)
if not self._first:
if self.register_watch_cb:
try:
utils.log(utils.cur(), "ready to callback")
self.register_watch_cb()
except Exception, err:
utils.err(utils.cur(), traceback.format_exc())
def safe_connect(self):
self._conn_cv.acquire()
self._handler = zookeeper.init(ZKConf.ZK_HOST, self.watcher_conn, ZKConf.ZOO_SESSION_TIMEOUT) # milliseconds
self._conn_cv.wait(8.0)
self._conn_cv.release()
try_times = 10
while not self._connected and try_times > 0:
self.connect()
try_times -= 1
if not self._connected:
raise ZooKeeperError("zookeeper [ %s ] safe_connect() failed" % str(ZKConf.ZK_HOST))
try:
user_pass = "******" % (ZKConf.ZOO_USERNAME, ZKConf.ZOO_PASSWORD)
utils.log(utils.cur(), self._name, "--add_auth--")
zookeeper.add_auth(self._handler, "digest", user_pass, None)
if self.register_node_cb:
try:
func = getattr(self, self.register_node_cb)
if func and callable(func):
utils.log(utils.cur(), self.register_node_cb, "%s callback %s" % ("*"*10, "*"*10))
def _do_register_node():
func(*self._args, **self._kwargs)
_do_register_node()
"""
check whether temporary nodes are still existed
when session timeout elapsed
because there are old temporary nodes when register new
temporary nodes
"""
threading.Timer(int(ZKConf.ZOO_SESSION_TIMEOUT/1000), _do_register_node).start()
except AttributeError, err:
utils.err(utils.cur(), traceback.format_exc())
utils.log(utils.cur(), self._first, self.register_watch_cb)
if not self._first:
if self.register_watch_cb:
try:
utils.log(utils.cur(), "ready to callback")
self.register_watch_cb()
except Exception, err:
utils.err(utils.cur(), traceback.format_exc())
def __init__(self, **config):
logger_name = config.get("logger_name")
self.logger = logging.getLogger(logger_name) if logger_name else Null()
self.zkhandle = None
self.auth = None
self.cv = threading.Condition()
try:
auth_config = config.get("auth")
if auth_config is not None:
auth_scheme = auth_config["scheme"]
auth_data = auth_config["data"]
self.auth = (auth_scheme, auth_data)
zklogfile_path, zklog_level = config.get("ZookeeperLog", ("/dev/stderr", "WARN"))
self.connection_timeout = config["timeout"]
self.zkhosts = ",".join(config["host"])
except KeyError as err:
self.logger.exception("Missing configuration option: %s", err)
raise
except Exception as err:
self.logger.exception("Unknown configuration error: %s", err)
raise
try:
_f = open(zklogfile_path, "a")
except IOError as err:
self.logger.error("Unable to open logfile %s %s", zklogfile_path, err)
else:
zookeeper.set_log_stream(_f)
zookeeper.set_debug_level(LOG_LEVELS.get(zklog_level.upper(), zookeeper.LOG_LEVEL_WARN))
self.connect()
if zookeeper.state(self.zkhandle) == zookeeper.CONNECTED_STATE:
self.logger.info("Connected to Zookeeper successfully")
else:
raise zookeeper.ZooKeeperException("Unable to connect " "to Zookeeper")
def on_auth_callback(state, result):
with self.cv:
if result == zookeeper.AUTHFAILED:
self.logger.error(zookeeper.zerror(zookeeper.AUTHFAILED))
self.logger.info("on_auth: state %s, result %s", state, result)
self.cv.notify()
if self.auth:
self.logger.info("Auth using %s", self.auth[0])
with self.cv:
res = zookeeper.add_auth(self.zkhandle, self.auth[0], self.auth[1], on_auth_callback)
if res != zookeeper.OK:
self.logger.error("Invalid status %d", zookeeper.zerror(res))
raise Exception("Invalid status")
self.cv.wait(self.connection_timeout)
if zookeeper.state(self.zkhandle) == zookeeper.AUTH_FAILED_STATE:
raise zookeeper.ZooKeeperException("authentication failed")
def add_auth(self, scheme, credentials):
'''Specifies the connection credentials
Overwrites zookeeper.add_auth() which is asynchronous
'''
class AuthWatch(object):
def __init__(self, condition):
self.condition = condition
self.has_fired = False
def __call__(self,handle, auth_result):
logger.debug('fired %d' % auth_result)
self.has_fired = True
if auth_result == zookeeper.OK:
self.condition.set()
logger.debug('Auth added')
else:
logger.warn('Could not set authentication')
# create condition
condition = threading.Event()
auth_watch = AuthWatch(condition)
# call method
logger.debug('Adding auth')
zookeeper.add_auth(self._handle, scheme, credentials, auth_watch)
# wait for completion
wait_time = self.recv_timeout()
logger.debug('waiting for %f seconds' % wait_time)
condition.wait(wait_time)
logger.debug('finished waiting')
if not auth_watch.has_fired:
logger.warn('Zookeeper server did not respond within %.2f seconds' % wait_time)
# return result
return condition.isSet()
def add_auth(self, scheme, cert):
"""
specify application credentials.
The application calls this function to specify its credentials for
purposes of authentication. The server will use the security provider
specified by the scheme parameter to authenticate the client
connection. If the authentication request has failed:
- the server connection is dropped
- the watcher is called with the AUTH_FAILED_STATE value as the state
parameter.
PARAMETERS:
@param scheme: the id of authentication scheme. Natively supported:
'digest' password-based authentication
@param cert: application credentials. The actual value depends on the
scheme.
RETURNS None.
If error occurs, one of the following corresponding exceptions will be
thrown.
OK on success or one of the following errcodes on failure:
AUTHFAILED authentication failed
BADARGUMENTS - invalid input parameters
INVALIDSTATE - zhandle state is either SESSION_EXPIRED_STATE or \
AUTH_FAILED_STATE
MARSHALLINGERROR - failed to marshall a request; possibly, out of \
memory
SYSTEMERROR - a system error occurred
"""
pc = utils.StatePipeCondition()
ok = zookeeper.add_auth(self._zhandle, scheme, cert,
functools.partial(_generic_completion,
pc))
assert ok == zookeeper.OK
results = pc.wait_and_get()
#unpack result as void_completion
handle, rc = results
assert handle == self._zhandle
if rc == zookeeper.OK:
return
self._raise_exception(rc)
def add_auth(self, scheme, identity):
"""Adds an authentication identity to this connection.
A connection can use multiple authentication identities at the
same time, all are checked when verifying acls on a node.
@param scheme: a string specifying a an authentication scheme
valid values include 'digest'.
@param identity: a string containing username and password colon
separated, for example 'mary:apples'
"""
d = defer.Deferred()
if self._check_connected(d):
return d
def _cb_authenticated(result_code):
if self._check_result(result_code, d):
return
d.callback(self)
callback = self._zk_thread_callback(_cb_authenticated)
result = zookeeper.add_auth(self.handle, scheme, identity, callback)
self._check_result(result, d)
return d
def delete(self, path): zookeeper.add_auth(self.handle, 'digest', AUTH, None) return zookeeper.delete(self.handle, path)
def set(self, path, value): zookeeper.add_auth(self.handle, 'digest', AUTH, None) return zookeeper.set(self.handle, path, value)
def create(self, path, value): zookeeper.add_auth(self.handle, 'digest', AUTH, None) return zookeeper.create(self.handle, path, value, ZOO_ACL)
def add_auth(self, scheme, data, callback):
return zookeeper.add_auth(self.handle, scheme, data, callback)
def delete(self, path):
zookeeper.add_auth(self.handle, 'digest', AUTH, None)
return zookeeper.delete(self.handle, path)
def set(self, path, value):
zookeeper.add_auth(self.handle, 'digest', AUTH, None)
return zookeeper.set(self.handle, path, value)
def create(self, path, value):
zookeeper.add_auth(self.handle, 'digest', AUTH, None)
return zookeeper.create(self.handle, path, value, ZOO_ACL)
def __init__(self, **config):
logger_name = config.get('logger_name')
self.logger = logging.getLogger(logger_name) if logger_name else Null()
self.zkhandle = None
self.auth = None
self.cv = threading.Condition()
try:
auth_config = config.get("auth")
if auth_config is not None:
auth_scheme = auth_config["scheme"]
auth_data = auth_config["data"]
self.auth = (auth_scheme, auth_data)
zklogfile_path, zklog_level = config.get("ZookeeperLog",
("/dev/stderr", "WARN"))
self.connection_timeout = config['timeout']
self.zkhosts = ','.join(config['host'])
except KeyError as err:
self.logger.exception("Missing configuration option: %s", err)
raise
except Exception as err:
self.logger.exception("Unknown configuration error: %s", err)
raise
try:
_f = open(zklogfile_path, 'a')
except IOError as err:
self.logger.error("Unable to open logfile %s %s",
zklogfile_path, err)
else:
zookeeper.set_log_stream(_f)
zookeeper.set_debug_level(LOG_LEVELS.get(zklog_level.upper(),
zookeeper.LOG_LEVEL_WARN))
self.connect()
if zookeeper.state(self.zkhandle) == zookeeper.CONNECTED_STATE:
self.logger.info('Connected to Zookeeper successfully')
else:
raise zookeeper.ZooKeeperException('Unable to connect '
'to Zookeeper')
def on_auth_callback(state, result):
with self.cv:
if result == zookeeper.AUTHFAILED:
self.logger.error(zookeeper.zerror(zookeeper.AUTHFAILED))
self.logger.info("on_auth: state %s, result %s",
state, result)
self.cv.notify()
if self.auth:
self.logger.info("Auth using %s", self.auth[0])
with self.cv:
res = zookeeper.add_auth(self.zkhandle, self.auth[0],
self.auth[1], on_auth_callback)
if res != zookeeper.OK:
self.logger.error("Invalid status %d",
zookeeper.zerror(res))
raise Exception("Invalid status")
self.cv.wait(self.connection_timeout)
if zookeeper.state(self.zkhandle) == zookeeper.AUTH_FAILED_STATE:
raise zookeeper.ZooKeeperException('authentication failed')
