예제 #1
0
def test_factory_create():
    factory = RedLockFactory([{"host": "localhost"}])

    lock = factory.create_lock("test_factory_create", ttl=500, retry_times=5, retry_delay=100)

    assert factory.redis_nodes == lock.redis_nodes
    assert factory.quorum == lock.quorum
    assert lock.ttl == 500
    assert lock.retry_times == 5
    assert lock.retry_delay == 100
    assert lock.factory == factory
예제 #2
0
def test_factory_create_from_url():
    factory = RedLockFactory([{"url": "redis://localhost/0"}])

    lock = factory.create_lock(
        "test_factory_create_from_url", ttl=500, retry_times=5, retry_delay=100
    )

    assert factory.redis_nodes == lock.redis_nodes
    assert factory.quorum == lock.quorum
    assert lock.ttl == 500
    assert lock.retry_times == 5
    assert lock.retry_delay == 100
    assert lock.factory == factory
예제 #3
0
    if parsed.scheme != 'redis':
        raise URLError('Redis URL does not have the redis scheme')

    path = parsed.path[1:] or ''
    query = parse_qs(parsed.query or '')
    if path:
        db = int(path)
    elif 'db' in query:
        db = int(query['db'][0])
    else:
        db = 0

    options = {'host': parsed.hostname, 'port': parsed.port, 'db': db}

    if parsed.password:
        options['password'] = parsed.password

    return options


if __name__ == '__main__':
    redis_url = 'redis://*****:*****@host:1234/?db=1'
    print(parse_redis_url(redis_url))
    exit(0)

from redlock import RedLockFactory
from rps_cnc import config

lock_factory = RedLockFactory(
    connection_details=[parse_redis_url(url) for url in config.REDIS_MASTERS])
예제 #4
0
class RedisCache(object):
    """
    A wrapper around the redis cache cluster which adds tenancy
    """
    conn = None
    tenant = None
    disabled = False
    redlock = None

    def __init__(self, tenant=None, service_name=None, redis_server=None):
        conn_info = current_app.config.get("redis_connection_info")
        if not redis_server:
            redis_server = current_app.config.get("redis_server", None)
        self.disabled = conn_info.get("disabled", False)
        if self.disabled:
            log.warning("Redis is disabled!")
            return
        if not redis_server:
            try:
                redis_server = g.driftenv_objects["redis_server"]
            except Exception:
                log.info(
                    "'redis_server' not found in config. Using default server '%s'",
                    conn_info["host"])
                redis_server = conn_info["host"]
        self.tenant = tenant or g.driftenv["name"]
        self.service_name = service_name or current_app.config["name"]
        self.host = redis_server
        self.port = conn_info["port"]
        self.conn = redis.StrictRedis(
            host=self.host,
            port=self.port,
            socket_timeout=conn_info.get("socket_timeout", 5),
            socket_connect_timeout=conn_info.get("socket_connect_timeout", 5),
            db=REDIS_DB,
        )

        self.key_prefix = "{}.{}:".format(self.tenant, self.service_name)

        self.redlock_factory = RedLockFactory(connection_details=[{
            'host':
            self.host,
            'port':
            self.port,
            'db':
            REDIS_DB,
        }], )

        log.debug("RedisCache initialized. self.conn = %s", self.conn)

    def make_key(self, key):
        """
        Create a redis key with tenant and service embedded into the key
        """
        return self.key_prefix + key

    def dump_object(self, value):
        """Dumps an object into a string for redis.  By default it serializes
        integers as regular string and pickle dumps everything else.
        """
        t = type(value)
        if t in integer_types:
            return str(value).encode('ascii')
        return b'!' + pickle.dumps(value)

    def load_object(self, value):
        """The reversal of :meth:`dump_object`.  This might be called with
        None.
        """
        if value is None:
            return None
        if value.startswith(b'!'):
            try:
                return pickle.loads(value[1:])
            except pickle.PickleError:
                return None
        try:
            return int(value)
        except ValueError:
            return value

    def set(self, key, value, expire=-1):
        compound_key = self.make_key(key)
        dump = self.dump_object(value)
        if expire == -1:
            result = self.conn.set(name=compound_key, value=dump)
        else:
            result = self.conn.setex(name=compound_key,
                                     value=dump,
                                     time=expire)

        return result

    def get(self, key):
        compound_key = self.make_key(key)
        try:
            ret = self.conn.get(compound_key)
        except redis.RedisError as e:
            self._log_error(e)
            return None

        ret = self.load_object(ret)
        return ret

    def delete(self, key):
        """
        Delete the item with the specified key
        """
        if self.disabled:
            log.info("Redis disabled. Not deleting key '%s'", key)
            return None
        compound_key = self.make_key(key)
        self.conn.delete(compound_key)

    def incr(self, key, amount=1, expire=None):
        """
        Increments the value of 'key' by 'amount'. If no key exists,
        the value will be initialized as 'amount'
        """
        if self.disabled:
            log.info("Redis disabled. Not incrementing key '%s'", key)
            return None
        compound_key = self.make_key(key)
        self.conn.incr(compound_key, amount)
        if expire:
            self.conn.expire(compound_key, expire)

    def lock(self, lock_name):
        return self.redlock_factory.create_lock(self.make_key(lock_name),
                                                retry_times=20,
                                                retry_delay=300)
예제 #5
0
    def setup_firewall(self, port, dnsport, nslist, family, subnets, udp,
                       user):
        if family not in [socket.AF_INET, socket.AF_INET6]:
            raise Exception(
                'Address family "%s" unsupported by tproxy method' %
                family_to_string(family))

        table = "mangle"

        def _ipt(*args):
            return ipt(family, table, *args)

        def _ipt_ttl(*args):
            return ipt_ttl(family, table, *args)

        def _ipt_proto_ports(proto, fport, lport):
            return proto + ('--dport', '%d:%d' % (fport, lport)) \
                    if fport else proto

        mark_chain = 'sshuttle-m-%s' % port
        tproxy_chain = 'sshuttle-t-%s' % port
        divert_chain = 'sshuttle-d-%s' % port

        # basic cleanup/setup of chains
        self.restore_firewall(port, family, udp, user)
        _ipt('-N', mark_chain)
        _ipt('-F', mark_chain)
        _ipt('-N', divert_chain)
        _ipt('-F', divert_chain)
        _ipt('-N', tproxy_chain)
        _ipt('-F', tproxy_chain)

        if (REDIS_HOST is None or REDIS_PORT is None):
            raise Fatal(
                "REDIS_HOST and REDIS_PORT environment variables must both be set!"
            )

        redlockFactory = RedLockFactory([{
            "host": REDIS_HOST,
            "port": REDIS_PORT
        }])
        lock = redlockFactory.create_lock(
            "SSHUTTLE_TPROXY_INSTANCE_CREATION_LOCK",
            ttl=500,
            retry_times=5,
            retry_delay=100)
        locked = lock.acquire()
        if locked == True:
            rule_count = ipt_rule_count(family, 'mangle', 'PREROUTING')
            if rule_count == 0:
                _ipt('-I', 'PREROUTING', '1', '-j', tproxy_chain)
            else:
                global iptables_lb_every
                iptables_lb_every = str(rule_count + 1)
                _ipt('-I', 'PREROUTING', '1', '-m', 'statistic', '--mode',
                     'nth', '--every', iptables_lb_every, '--packet', '0',
                     '-j', tproxy_chain)
            lock.release()
        else:
            lock.release()
            raise Exception('Failed to acquire lock to edit iptable')

        _ipt('-I', 'OUTPUT', '1', '-j', mark_chain)
        _ipt('-A', divert_chain, '-j', 'MARK', '--set-mark', '1')
        _ipt('-A', divert_chain, '-j', 'ACCEPT')
        _ipt('-A', tproxy_chain, '-m', 'socket', '-j', divert_chain, '-m',
             'tcp', '-p', 'tcp')

        # get the address of eth0
        ni.ifaddresses('eth0')
        ip = ni.ifaddresses('eth0')[2][0]['addr']

        # add a rule not route packets that are
        # generated locally though sshuttle
        _ipt('-A', mark_chain, '-j', 'RETURN', '--src', '%s/32' % ip)

        # add a rule to not route packets that are
        # destined to the local address though sshuttle
        _ipt('-A', mark_chain, '-j', 'RETURN', '--dest', '%s/32' % ip)

        if udp:
            _ipt('-A', tproxy_chain, '-m', 'socket', '-j', divert_chain, '-m',
                 'udp', '-p', 'udp')

        for _, ip in [i for i in nslist if i[0] == family]:
            _ipt('-A', mark_chain, '-j', 'MARK', '--set-mark', '1', '--dest',
                 '%s/32' % ip, '-m', 'udp', '-p', 'udp', '--dport', '53')
            _ipt('-A', tproxy_chain, '-j', 'TPROXY', '--tproxy-mark',
                 '0x1/0x1', '--dest', '%s/32' % ip, '-m', 'udp', '-p', 'udp',
                 '--dport', '53', '--on-port', str(dnsport))

        for _, swidth, sexclude, snet, fport, lport \
                in sorted(subnets, key=subnet_weight, reverse=True):
            tcp_ports = ('-p', 'tcp')
            tcp_ports = _ipt_proto_ports(tcp_ports, fport, lport)

            if sexclude:
                _ipt('-A', mark_chain, '-j', 'RETURN', '--dest',
                     '%s/%s' % (snet, swidth), '-m', 'tcp', *tcp_ports)
                _ipt('-A', tproxy_chain, '-j', 'RETURN', '--dest',
                     '%s/%s' % (snet, swidth), '-m', 'tcp', *tcp_ports)
            else:
                _ipt('-A', mark_chain, '-j', 'MARK', '--set-mark', '1',
                     '--dest', '%s/%s' % (snet, swidth), '-m', 'tcp',
                     *tcp_ports)
                _ipt('-A', tproxy_chain, '-j', 'TPROXY', '--tproxy-mark',
                     '0x1/0x1', '--dest', '%s/%s' % (snet, swidth), '-m',
                     'tcp', *(tcp_ports + ('--on-port', str(port))))

            if udp:
                udp_ports = ('-p', 'udp')
                udp_ports = _ipt_proto_ports(udp_ports, fport, lport)

                if sexclude:
                    _ipt('-A', mark_chain, '-j', 'RETURN', '--dest',
                         '%s/%s' % (snet, swidth), '-m', 'udp', *udp_ports)
                    _ipt('-A', tproxy_chain, '-j', 'RETURN', '--dest',
                         '%s/%s' % (snet, swidth), '-m', 'udp', *udp_ports)
                else:
                    _ipt('-A', mark_chain, '-j', 'MARK', '--set-mark', '1',
                         '--dest', '%s/%s' % (snet, swidth), '-m', 'udp', '-p',
                         'udp')
                    _ipt('-A', tproxy_chain, '-j', 'TPROXY', '--tproxy-mark',
                         '0x1/0x1', '--dest', '%s/%s' % (snet, swidth), '-m',
                         'udp', *(udp_ports + ('--on-port', str(port))))
예제 #6
0
def make_lock_factory(connection_details):
    return RedLockFactory(connection_details)
예제 #7
0
파일: redis.py 프로젝트: dgnorth/drift
class RedisCache(object):
    """
    A wrapper around the redis cache cluster which adds tenancy
    """
    conn = None
    tenant = None
    disabled = False
    redlock = None

    def __init__(self, tenant, service_name, redis_config):
        self.disabled = redis_config.get("disabled", False)
        if self.disabled:
            log.warning("Redis is disabled!")
            return

        self.tenant = tenant
        self.service_name = service_name
        self.host = redis_config["host"]
        self.port = redis_config["port"]

        # Override Redis hostname if needed
        if os.environ.get('DRIFT_USE_LOCAL_SERVERS', False):
            self.host = 'localhost'

        self.conn = redis.StrictRedis(
            host=self.host,
            port=self.port,
            socket_timeout=redis_config.get("socket_timeout", 5),
            socket_connect_timeout=redis_config.get("socket_connect_timeout",
                                                    5),
            db=redis_config.get("db_number", REDIS_DB),
        )

        self.key_prefix = "{}.{}:".format(self.tenant, self.service_name)

        self.redlock_factory = RedLockFactory(connection_details=[{
            'host':
            self.host,
            'port':
            self.port,
            'db':
            redis_config.get("db_number", REDIS_DB),
        }], )

        log.debug("RedisCache initialized. self.conn = %s", self.conn)

    def make_key(self, key):
        """
        Create a redis key with tenant and service embedded into the key
        """
        return self.key_prefix + key

    def dump_object(self, value):
        """Dumps an object into a string for redis.  By default it serializes
        integers as regular string and pickle dumps everything else.
        """
        t = type(value)
        if t in integer_types:
            return str(value).encode('ascii')
        return b'!' + pickle.dumps(value)

    def load_object(self, value):
        """The reversal of :meth:`dump_object`.  This might be called with
        None.
        """
        if value is None:
            return None
        if value.startswith(b'!'):
            try:
                return pickle.loads(value[1:])
            except pickle.PickleError:
                return None
        try:
            return int(value)
        except ValueError:
            return value

    def set(self, key, value, expire=-1):
        compound_key = self.make_key(key)
        dump = self.dump_object(value)
        if expire == -1:
            result = self.conn.set(name=compound_key, value=dump)
        else:
            result = self.conn.setex(name=compound_key,
                                     value=dump,
                                     time=expire)

        return result

    def get(self, key):
        compound_key = self.make_key(key)
        try:
            ret = self.conn.get(compound_key)
        except redis.RedisError as e:
            log.exception("Can't fetch key '%s'", compound_key)
            return None

        ret = self.load_object(ret)
        return ret

    def delete(self, key):
        """
        Delete the item with the specified key
        """
        if self.disabled:
            log.info("Redis disabled. Not deleting key '%s'", key)
            return None
        compound_key = self.make_key(key)
        self.conn.delete(compound_key)

    def incr(self, key, amount=1, expire=None):
        """
        Increments the value of 'key' by 'amount'. If no key exists,
        the value will be initialized as 'amount'
        """
        if self.disabled:
            log.info("Redis disabled. Not incrementing key '%s'", key)
            return None
        compound_key = self.make_key(key)
        self.conn.incr(compound_key, amount)
        if expire:
            self.conn.expire(compound_key, expire)

    def lock(self, lock_name):
        return self.redlock_factory.create_lock(self.make_key(lock_name),
                                                retry_times=20,
                                                retry_delay=300)

    def delete_all(self):
        """remove all the keys for this tenant from redis"""
        if self.disabled:
            log.info("Redis disabled. Not deleting all keys")
            return
        compound_key = self.make_key("*")
        for key in self.conn.scan_iter(compound_key):
            self.conn.delete(key)