class Redis:
def __init__(self, app, user_id=0):
self.redis_client = FlaskRedis(app)
self.key = 'ctfd_whale_lock-' + str(user_id)
self.lock = None
self.global_port_key = "ctfd_whale-port-set"
self.global_network_key = "ctfd_whale-network-set"
def init_redis_port_sets(self):
configs = constant.frp_api()
self.redis_client.delete(self.global_port_key)
self.redis_client.delete(self.global_network_key)
containers = DBUtils.get_all_container()
used_port_list = []
for container in containers:
if container.port != 0:
used_port_list.append(container.port)
# for port in range(int(configs.get("frp_direct_port_minimum", 29000)), int(configs.get("frp_direct_port_maximum", 28000)) + 1):
for port in range(int(configs.frp_direct_port_minimum),
int(configs.frp_direct_port_maximum)):
if port not in used_port_list:
self.add_available_port(port)
def add_available_port(self, port):
self.redis_client.sadd(self.global_port_key, str(port))
def get_available_port(self):
return int(self.redis_client.spop(self.global_port_key))
def acquire_lock(self):
lock = self.redis_client.lock(name=self.key, timeout=20)
if not lock.acquire(blocking=True, blocking_timeout=2.0):
return False
self.lock = lock
return True
def release_lock(self):
if self.lock is None:
return False
try:
self.lock.release()
return True
except LockError:
return False
class KSVotesRedis():
def __init__(self, app=current_app):
self.redis_client = FlaskRedis(app)
self.namespace = os.getenv('APP_CONFIG') or 'default'
def get(self, key):
return self.redis_client.get(self.namespace + ":" + key)
def set(self, key, value, ttl=60):
return self.redis_client.set(self.namespace + ":" + key, value, ex=ttl)
def clear(self, key):
return self.redis_client.delete(self.namespace + ":" + key)
def get_or_set(self, key, setter, ttl):
ns_key = self.namespace + ":" + key
with self.redis_client.pipeline() as pipe:
try:
pipe.watch(ns_key)
# after WATCHing, the pipeline is put into immediate execution
# mode until we tell it to start buffering commands again.
# this allows us to get the current value of our sequence
if pipe.exists(ns_key):
return pipe.get(ns_key)
# now we can put the pipeline back into buffered mode with MULTI
pipe.multi()
pipe.set(ns_key, setter(), ex=ttl)
pipe.get(ns_key)
# and finally, execute the pipeline (the set and get commands)
return pipe.execute()[-1]
# if a WatchError wasn't raised during execution, everything
# we just did happened atomically.
except WatchError:
# another client must have changed key between
# the time we started WATCHing it and the pipeline's execution.
# Let's just get the value they changed it to.
return pipe.get(ns_key)
class RedisDB(object):
"""Simple Queue with Redis Backend"""
default_namespace = None
def __init__(self,
app=None,
name=None,
namespace=None,
sub_namespace=None):
'key format -> namespace:sub_namescape:name'
self.name = name
self.namespace = namespace or self.default_namespace
self.sub_namespace = sub_namespace
if app is not None:
self.init(app)
def init(self, app):
self.app = app
self.app.config.setdefault('REDIS_QUEUE_NAMESPACE', self.namespace)
self.app.config.setdefault('REDIS_QUEUE_NAME', self.name)
self.db = FlaskRedis(app)
self.namespace = self.namespace or self.app.config.get(
'REDIS_QUEUE_NAMESPACE')
self.name = self.name or self.app.config.get('REDIS_QUEUE_NAME')
def qsize(self, name=None, sub_namespace=None):
"""Return the approximate size of the queue."""
return self.db.llen(self.abs_key(name, sub_namespace))
def empty(self, name):
"""Return True if the queue is empty, False otherwise."""
return self.qsize(name) == 0
def put(self, item, name, sub_namespace=None):
"""Put item into the queue."""
item = self.jsonfiy_value(item)
self.db.lpush(self.abs_key(name, sub_namespace), item)
def get(self, name=None, block=True, timeout=None):
"""Remove and return an item from the queue.
If optional args block is true and timeout is None (the default), block
if necessary until an item is available."""
key = '%s:%s' % (self.namespace, name or self.name)
if block:
item = self.db.blpop(key, timeout=timeout)
else:
item = self.db.lpop(key)
if item:
item = item[1]
return item
def get_nowait(self, name):
"""Equivalent to get(False)."""
return self.get(name, False)
def get_list(self, name, start=0, end=-1):
key = '%s:%s' % (self.namespace, name or self.name)
return self.db.lrange(key, start, end)
def db(self):
return self.db
def delete_queue(self, *names):
if names:
keys = ["%s:%s" % (self.namespace, name) for name in names]
self.db.delete(*keys)
else:
key = self.abs_key(self.name)
self.db.delete(key)
def log(self, value):
self.db.lpush('log', '%s: %s' % (str(datetime.now())[:19], value))
def setex(self, name, value, sub_namespace=None, time=300):
value = self.jsonfiy_value(value)
self.db.setex(self.abs_key(name, sub_namespace), value, time)
def hset(self, field, value, name, sub_namespace=None):
value = self.jsonfiy_value(value)
self.db.hset(self.abs_key(name, sub_namespace), field, value)
def hlen(self, name, sub_namespace=None):
return self.db.hlen(self.abs_key(name, sub_namespace))
def hkeys(self, name, sub_namespace=None):
return self.db.hkeys(self.abs_key(name, sub_namespace))
def hget(self, name, field, sub_namespace=None):
value = self.db.hget(self.abs_key(name, sub_namespace), field)
return json.loads(value) if value else None
def abs_key(self, name, sub_namespace=None):
name = name or self.name
sub_namespace = sub_namespace or self.sub_namespace
return "%s:%s" % (self.namespace, "%s:%s" %
(sub_namespace, name) if sub_namespace else name)
def jsonfiy_value(self, value):
if isinstance(value, dict):
value = json.dumps(value)
return value
def expire(self, name, seconds=300, sub_namespace=None):
self.db.expire(self.abs_key(name, sub_namespace), seconds)
def set(self, name, value, sub_namespace=None):
value = self.jsonfiy_value(value)
self.db.set(self.abs_key(name, sub_namespace), value)
def rpush(self, item, name=None, sub_namespace=None):
item = self.jsonfiy_value(item)
self.db.rpush(self.abs_key(name, sub_namespace), item)
def exists(self, name=None, sub_namespace=None):
return self.db.exists(self.abs_key(name, sub_namespace))
def lrange(self, start, stop, name=None, sub_namespace=None):
return self.db.lrange(self.abs_key(name, sub_namespace), start, stop)
def sadd(self, name=None, sub_namespace=None, *value):
return self.db.sadd(self.abs_key(name, sub_namespace), value)
def zincrby(self, member, increment=1, name=None, sub_namespace=None):
return self.db.zincrby(self.abs_key(name, sub_namespace), member,
increment)
def zrange(self, name, start=0, stop=-1):
return self.db.zrange(self.abs_key(name), start, stop)
def zscore(self, name, member):
return self.db.zscore(self.abs_key(name), member)
def incr(self, name=None, sub_namespace=None):
return self.db.incr(self.abs_key(name, sub_namespace))
def delete(self, name=None, sub_namespace=None):
return self.db.delete(self.abs_key(name, sub_namespace))
class RedisUtils:
def __init__(self, app, user_id=0):
self.redis_client = FlaskRedis(app)
self.key = 'ctfd_whale_lock-' + str(user_id)
self.lock = None
self.global_port_key = "ctfd_whale-port-set"
self.global_network_key = "ctfd_whale-network-set"
def init_redis_port_sets(self):
configs = DBUtils.get_all_configs()
self.redis_client.delete(self.global_port_key)
self.redis_client.delete(self.global_network_key)
containers = DBUtils.get_all_container()
used_port_list = []
for container in containers:
if container.port != 0:
used_port_list.append(container.port)
for port in range(int(configs.get("frp_direct_port_minimum", 29000)), int(configs.get("frp_direct_port_maximum", 28000)) + 1):
if port not in used_port_list:
self.add_available_port(port)
client = docker.DockerClient(base_url=configs.get("docker_api_url"))
docker_subnet = configs.get("docker_subnet", "174.1.0.0/16")
try:
docker_subnet = unicode(docker_subnet)
except:
pass
docker_subnet_new_prefix = int(configs.get("docker_subnet_new_prefix", "24"))
exist_networks = []
available_networks = []
for network in client.networks.list(filters={'label': 'prefix'}):
exist_networks.append(str(network.attrs['Labels']['prefix']))
for network in list(ipaddress.ip_network(docker_subnet).subnets(new_prefix=docker_subnet_new_prefix)):
if str(network) not in exist_networks:
available_networks.append(str(network))
self.redis_client.sadd(self.global_network_key, *set(available_networks))
def add_available_network_range(self, network_range):
self.redis_client.sadd(self.global_network_key, network_range.encode())
def get_available_network_range(self):
return self.redis_client.spop(self.global_network_key).decode()
def add_available_port(self, port):
self.redis_client.sadd(self.global_port_key, str(port))
def get_available_port(self):
return int(self.redis_client.spop(self.global_port_key))
def acquire_lock(self):
lock = self.redis_client.lock(name=self.key, timeout=10)
if not lock.acquire(blocking=True, blocking_timeout=2.0):
return False
self.lock = lock
return True
def release_lock(self):
if self.lock is None:
return False
try:
self.lock.release()
return True
except LockError:
return False
