AIL-framework/bin/core/ail_2_ail.py

#!/usr/bin/env python3
# -*-coding:UTF-8 -*

import os
import json
import secrets
import sys
import time
import uuid

sys.path.append(os.path.join(os.environ['AIL_BIN'], 'lib/'))
import ConfigLoader

sys.path.append(os.path.join(os.environ['AIL_BIN'], 'core/'))
import screen

sys.path.append(os.path.join(os.environ['AIL_BIN'], 'packages/'))
from Item import Item

config_loader = ConfigLoader.ConfigLoader()
r_cache = config_loader.get_redis_conn("Redis_Cache")
r_serv_db = config_loader.get_redis_conn("ARDB_DB")
r_serv_sync = config_loader.get_redis_conn("ARDB_DB")
config_loader = None

def generate_uuid():
    return str(uuid.uuid4()).replace('-', '')

def generate_sync_api_key():
    return secrets.token_urlsafe(42)

def get_ail_uuid():
    return r_serv_db.get('ail:uuid')

# # TODO: # TODO: # TODO: # TODO: # TODO: ADD SYNC MODE == PUSH
# # TODO: get connection status
# # TODO: get connection METADATA
#############################
#                           #
#### SYNC CLIENT MANAGER ####

def get_all_sync_clients(r_set=False):
    res = r_cache.smembers('ail_2_ail:all_sync_clients')
    if r_set:
        return set(res)
    else:
        return res

def get_sync_client_ail_uuid(client_id):
    return r_cache.hget(f'ail_2_ail:sync_client:{client_id}', 'ail_uuid')

def get_sync_client_queue_uuid(client_id):
    return r_cache.hget(f'ail_2_ail:sync_client:{client_id}', 'queue_uuid')

def delete_sync_client_cache(client_id):
    ail_uuid = get_sync_client_ail_uuid(client_id)
    queue_uuid = get_sync_client_queue_uuid(client_id)
    # map ail_uuid/queue_uuid
    r_cache.srem(f'ail_2_ail:ail_uuid:{ail_uuid}', client_id)
    r_cache.srem(f'ail_2_ail:queue_uuid:{queue_uuid}', client_id)

    r_cache.delete(f'ail_2_ail:sync_client:{client_id}')
    r_cache.srem('ail_2_ail:all_sync_clients', client_id)

def delete_all_sync_clients_cache():
    for client_id in get_all_sync_clients():
        delete_sync_client_cache(client_id)
    r_cache.delete('ail_2_ail:all_sync_clients')

# command: -launch
#          -kill
#          -relaunch
## TODO: check command
def send_command_to_manager(command, client_id=-1):
    dict_action = {'command': command, 'client_id': client_id}
    r_cache.sadd('ail_2_ail:client_manager:command', dict_action)

class AIL2AILClientManager(object):
    """AIL2AILClientManager."""

    SCREEN_NAME = 'AIL_2_AIL'
    SCRIPT_NAME = 'ail_2_ail_client.py'
    SCRIPT_DIR = os.path.join(os.environ['AIL_BIN'], 'core')

    def __init__(self):
        # dict client_id: AIL2AILCLIENT or websocket
        self.clients = {}
        # launch all sync clients
        self.relaunch_all_sync_clients()

    def get_all_clients(self):
        return self.clients

    # return new client id
    def get_new_sync_client_id(self):
        for new_id in range(100000):
            if new_id not in self.clients:
                return str(new_id)

    def get_sync_client_ail_uuid(self, client_id):
        return self.clients[client_id]['ail_uuid']

    def get_sync_client_queue_uuid(self, client_id):
        return self.clients[client_id]['queue_uuid']

    # # TODO: check PUSH ACL
    def get_all_sync_clients_to_launch(self):
        return get_all_ail_instance()

    def relaunch_all_sync_clients(self):
        delete_all_sync_clients_cache()
        self.clients = {}
        for ail_uuid in self.get_all_sync_clients_to_launch():
             self.launch_sync_client(ail_uuid)

    def launch_sync_client(self, ail_uuid):
        dir_project = os.environ['AIL_HOME']
        client_id = self.get_new_sync_client_id()
        script_options = f'-a {ail_uuid} -m push -i {client_id}'
        screen.create_screen(AIL2AILClientManager.SCREEN_NAME)
        screen.launch_uniq_windows_script(AIL2AILClientManager.SCREEN_NAME,
                                            client_id, dir_project,
                                            AIL2AILClientManager.SCRIPT_DIR,
                                            AIL2AILClientManager.SCRIPT_NAME,
                                            script_options=script_options, kill_previous_windows=True)
        # save sync client status
        r_cache.hset(f'ail_2_ail:sync_client:{client_id}', 'ail_uuid', ail_uuid)
        r_cache.hset(f'ail_2_ail:sync_client:{client_id}', 'launch_time', int(time.time()))

        # create map ail_uuid/queue_uuid
        r_cache.sadd(f'ail_2_ail:ail_uuid:{ail_uuid}', client_id)

        self.clients[client_id] = {'ail_uuid': ail_uuid}

    # # TODO: FORCE KILL ????????????
    # # TODO: check if exists
    def kill_sync_client(self, client_id):
        if not kill_screen_window('AIL_2_AIL', client_id):
            # # TODO: log kill error
            pass

        delete_sync_client_cache(client_id)
        self.clients.pop(client_id)

    ## COMMANDS ##

    def get_manager_command(self):
        res = r_cache.spop('ail_2_ail:client_manager:command')
        if res:
            return json.dumps(res)
        else:
            return None

    def execute_manager_command(self, command_dict):
        command = command_dict.get('command')
        if command == 'launch':
            ail_uuid = int(command_dict.get('ail_uuid'))
            queue_uuid = int(command_dict.get('queue_uuid'))
            self.launch_sync_client(ail_uuid, queue_uuid)
        elif command == 'relaunch':
            self.relaunch_all_sync_clients()
        else:
            # only one sync client
            client_id = int(command_dict.get('client_id'))
            if client_id < 1:
                print('Invalid client id')
                return None
            if command == 'kill':
                self.kill_sync_client(client_id)
            elif command == 'relaunch':
                ail_uuid = self.get_sync_client_ail_uuid(client_id)
                queue_uuid = self.get_sync_client_queue_uuid(client_id)
                self.kill_sync_client(client_id)
                self.launch_sync_client(ail_uuid, queue_uuid)

########################################
########################################
########################################

# # TODO: ADD METADATA
def get_sync_client_status(client_id):
    dict_client = {'id': client_id}
    dict_client['ail_uuid'] = get_sync_client_ail_uuid(client_id)
    dict_client['queue_uuid'] = get_sync_client_queue_uuid(client_id)
    return dict_client

def get_all_sync_client_status():
    sync_clients = []
    all_sync_clients = r_cache.smembers('ail_2_ail:all_sync_clients')
    for client_id in all_sync_clients:
        sync_clients.append(get_sync_client_status(client_id))
    return sync_clients

######################
#                    #
#### AIL INSTANCE ####

## AIL KEYS ##

def get_all_ail_instance_keys():
    return r_serv_sync.smembers(f'ail:instance:key:all')

def is_allowed_ail_instance_key(key):
    return r_serv_sync.sismember(f'ail:instance:key:all', key)

def get_ail_instance_key(ail_uuid):
    return r_serv_sync.hget(f'ail:instance:{ail_uuid}', 'api_key')

def get_ail_instance_by_key(key):
    return r_serv_sync.get(f'ail:instance:key:{key}')

# def check_acl_sync_queue_ail(ail_uuid, queue_uuid, key):
#     return is_ail_instance_queue(ail_uuid, queue_uuid)

def update_ail_instance_key(ail_uuid, new_key):
    old_key = get_ail_instance_key(ail_uuid)
    r_serv_sync.srem(f'ail:instance:key:all', old_key)
    r_serv_sync.delete(f'ail:instance:key:{old_key}')

    r_serv_sync.sadd(f'ail:instance:key:all', new_key)
    r_serv_sync.delete(f'ail:instance:key:{new_key}', ail_uuid)
    r_serv_sync.hset(f'ail:instance:{ail_uuid}', 'api_key', new_key)

#- AIL KEYS -#

def get_all_ail_instance():
    return r_serv_sync.smembers('ail:instance:all')

def get_ail_instance_all_sync_queue(ail_uuid):
    return r_serv_sync.smembers(f'ail:instance:sync_queue:{ail_uuid}')

def is_ail_instance_queue(ail_uuid, queue_uuid):
    return r_serv_sync.sismember(f'ail:instance:sync_queue:{ail_uuid}', queue_uuid)

def get_ail_instance_url(ail_uuid):
    return r_serv_sync.hget(f'ail:instance:{ail_uuid}', 'url')

def get_ail_instance_description(ail_uuid):
    return r_serv_sync.hget(f'ail:instance:{ail_uuid}', 'description')

def is_ail_instance_push_enabled(ail_uuid):
    res = r_serv_sync.hget(f'ail:instance:{ail_uuid}', 'push')
    return res == 'True'

def is_ail_instance_pull_enabled(ail_uuid):
    res = r_serv_sync.hget(f'ail:instance:{ail_uuid}', 'pull')
    return res == 'True'

def is_ail_instance_sync_enabled(ail_uuid, sync_mode=None):
    if sync_mode is None:
        return is_ail_instance_push_enabled(ail_uuid) or is_ail_instance_pull_enabled(ail_uuid)
    elif sync_mode == 'pull':
        return is_ail_instance_pull_enabled(ail_uuid)
    elif sync_mode == 'push':
        return is_ail_instance_push_enabled(ail_uuid)
    else:
        return False

def change_pull_push_state(ail_uuid, pull=False, push=False):
    # sanityze pull/push
    if pull:
        pull = True
    else:
        pull = False
    if push:
        push = True
    else:
        push = False
    r_serv_sync.hset(f'ail:instance:{ail_uuid}', 'push', push)
    r_serv_sync.hset(f'ail:instance:{ail_uuid}', 'pull', pull)


# # TODO: HIDE ADD GLOBAL FILTER (ON BOTH SIDE)
# # TODO: push/pull
def get_ail_instance_metadata(ail_uuid):
    dict_meta = {}
    dict_meta['uuid'] = ail_uuid
    dict_meta['url'] = get_ail_instance_url(ail_uuid)
    dict_meta['description'] = get_ail_instance_description(ail_uuid)

    # # TODO: HIDE
    dict_meta['api_key'] = get_ail_instance_key(ail_uuid)

    # # TODO:
    # - set UUID sync_queue

    return dict_meta

# # TODO: VALIDATE URL
#                  API KEY
def create_ail_instance(ail_uuid, url, api_key=None, description=None):
    r_serv_sync.sadd('ail:instance:all', ail_uuid)
    r_serv_sync.hset(f'ail:instance:{ail_uuid}', 'url', url)
    ## API KEY ##
    if not api_key:
        api_key = generate_sync_api_key()
    r_serv_sync.hset(f'ail:instance:{ail_uuid}', 'api_key', api_key)
    r_serv_sync.sadd('ail:instance:key:all', api_key)
    r_serv_sync.set(f'ail:instance:key:{api_key}', ail_uuid)
    #- API KEY -#
    if description:
        r_serv_sync.hset(f'ail:instance:{ail_uuid}', 'description', description)
    return ail_uuid

def delete_ail_instance(ail_uuid):
    for queue_uuid in get_ail_instance_all_sync_queue(ail_uuid):
        unregister_ail_to_sync_queue(ail_uuid, queue_uuid)
    r_serv_sync.delete(f'ail:instance:sync_queue:{ail_uuid}')
    key = get_ail_instance_by_key(ail_uuid)
    r_serv_sync.delete(f'ail:instance:{ail_uuid}')
    r_serv_sync.srem('ail:instance:key:all', ail_uuid)
    r_serv_sync.delete(f'ail:instance:key:{key}', ail_uuid)
    r_serv_sync.srem('ail:instance:all', ail_uuid)

def get_last_updated_ail_instance():
    epoch = r_serv_sync.get(f'ail:instance:queue:last_updated')
    if not epoch:
        epoch = 0
    return float(epoch)

####################
#                  #
#### SYNC QUEUE ####

class Sync_Queue(object): # # TODO: use for edit
    """Sync_Queue."""

    def __init__(self, uuid):
        self.uuid = uuid

def get_all_sync_queue():
    return r_serv_sync.smembers('ail2ail:sync_queue:all')

def get_sync_queue_all_ail_instance(queue_uuid):
    return r_serv_sync.smembers(f'ail2ail:sync_queue:ail_instance:{queue_uuid}')

# # TODO: check if push or pull enabled ?
def is_queue_used_by_ail_instace(queue_uuid):
    return r_serv_sync.exists(f'ail2ail:sync_queue:ail_instance:{queue_uuid}')

# # TODO: add others filter
def get_sync_queue_filter(queue_uuid):
    return r_serv_sync.smembers(f'ail2ail:sync_queue:filter:tags:{queue_uuid}')

# # TODO: ADD FILTER
def get_sync_queue_metadata(queue_uuid):
    dict_meta = {}
    dict_meta['uuid'] = queue_uuid
    dict_meta['name'] = r_serv_sync.hget(f'ail2ail:sync_queue:{queue_uuid}', 'name')
    dict_meta['description'] = r_serv_sync.hget(f'ail2ail:sync_queue:{queue_uuid}', 'description')
    dict_meta['max_size'] = r_serv_sync.hget(f'ail2ail:sync_queue:{queue_uuid}', 'max_size')

    # # TODO: TO ADD:
    # - set uuid instance

    return dict_meta

#####################################################
def get_all_sync_queue_dict():
    dict_sync_queues = {}
    for queue_uuid in get_all_sync_queue():
        if is_queue_used_by_ail_instace(queue_uuid):
            dict_queue = {}
            dict_queue['filter'] = get_sync_queue_filter(queue_uuid)

            dict_queue['ail_instances'] = [] ############ USE DICT ?????????
            for ail_uuid in get_sync_queue_all_ail_instance(queue_uuid):
                dict_ail = {'ail_uuid': ail_uuid,
                            'pull': is_ail_instance_pull_enabled(ail_uuid),
                            'push': is_ail_instance_push_enabled(ail_uuid)}
                if dict_ail['pull'] or dict_ail['push']:
                    dict_queue['ail_instances'].append(dict_ail)
            if dict_queue['ail_instances']:
                dict_sync_queues[queue_uuid] = dict_queue
    return dict_sync_queues


def register_ail_to_sync_queue(ail_uuid, queue_uuid):
    r_serv_sync.sadd(f'ail2ail:sync_queue:ail_instance:{queue_uuid}', ail_uuid)
    r_serv_sync.sadd(f'ail:instance:sync_queue:{ail_uuid}', queue_uuid)

def unregister_ail_to_sync_queue(ail_uuid, queue_uuid):
    r_serv_sync.srem(f'ail2ail:sync_queue:ail_instance:{queue_uuid}', ail_uuid)
    r_serv_sync.srem(f'ail:instance:sync_queue:{ail_uuid}', queue_uuid)

# # TODO: optionnal name ???
# # TODO: SANITYZE TAGS
def create_sync_queue(name, tags=[], description=None, max_size=100):
    queue_uuid = generate_uuid()
    r_serv_sync.sadd('ail2ail:sync_queue:all', queue_uuid)

    r_serv_sync.hset(f'ail2ail:sync_queue:{queue_uuid}', 'name', name)
    if description:
        r_serv_sync.hset(f'ail2ail:sync_queue:{queue_uuid}', 'description', description)
    r_serv_sync.hset(f'ail2ail:sync_queue:{queue_uuid}', 'max_size', max_size)

    for tag in tags:
        r_serv_sync.sadd(f'ail2ail:sync_queue:filter:tags:{queue_uuid}', tag)

    return queue_uuid

def delete_sync_queue(queue_uuid):
    for ail_uuid in get_sync_queue_all_ail_instance(queue_uuid):
        unregister_ail_to_sync_queue(ail_uuid, queue_uuid)
    r_serv_sync.delete(f'ail2ail:sync_queue:{queue_uuid}')
    r_serv_sync.srem('ail2ail:sync_queue:all', queue_uuid)
    return queue_uuid

#############################
#                           #
#### SYNC REDIS QUEUE #######

def get_sync_queue_object(ail_uuid, push=True):
    for queue_uuid in get_ail_instance_all_sync_queue(ail_uuid):
        obj_dict = get_sync_queue_object_by_queue_uuid(queue_uuid, ail_uuid, push=push)
        if obj_dict:
            return obj_dict
    return None

def get_sync_queue_object_by_queue_uuid(queue_uuid, ail_uuid, push=True):
    if push:
        sync_mode = 'push'
    else:
        sync_mode = 'pull'
    obj_dict = r_serv_sync.lpop(f'sync:queue:{sync_mode}:{queue_uuid}:{ail_uuid}')
    if obj_dict:
        obj_dict = json.loads(obj_dict)
        # # REVIEW: # TODO: create by obj type
        return Item(obj_dict['id'])

def add_object_to_sync_queue(queue_uuid, ail_uuid, obj_dict, push=True, pull=True):
    obj = json.dumps(obj_dict)

    # # TODO: # FIXME: USE CACHE ??????
    if push:
        r_serv_sync.lpush(f'sync:queue:push:{queue_uuid}:{ail_uuid}', obj)
        r_serv_sync.ltrim(f'sync:queue:push:{queue_uuid}:{ail_uuid}', 0, 200)

    if pull:
        r_serv_sync.lpush(f'sync:queue:pull:{queue_uuid}:{ail_uuid}', obj)
        r_serv_sync.ltrim(f'sync:queue:pull:{queue_uuid}:{ail_uuid}', 0, 200)

# # TODO: # REVIEW: USE CACHE ????? USE QUEUE FACTORY ?????
def get_sync_importer_ail_stream():
    return r_serv_sync.spop('sync:queue:importer')

def add_ail_stream_to_sync_importer(ail_stream):
    ail_stream = json.dumps(ail_stream)
    r_serv_sync.sadd('sync:queue:importer', ail_stream)

#############################
#                           #
#### AIL EXCHANGE FORMAT ####

def create_ail_stream(Object):
    ail_stream = {'format': 'ail',
                  'version': 1,
                  'type': Object.get_type()}

    # OBJECT META
    ail_stream['meta'] = {'ail_mime-type': 'text/plain'}
    ail_stream['meta']['ail:id'] = Object.get_id()
    ail_stream['meta']['ail:tags'] = Object.get_tags()
    # GLOBAL PAYLOAD
    ail_stream['meta']['ail:uuid'] = get_ail_uuid()

    # OBJECT PAYLOAD
    ail_stream['payload'] = Object.get_ail_2_ail_payload()

    return ail_stream

if __name__ == '__main__':

    ail_uuid = '03c51929-eeab-4d47-9dc0-c667f94c7d2d'
    url = "wss://localhost:4443"
    api_key = 'secret'
    #description = 'first test instance'
    queue_uuid = '79bcafc0a6d644deb2c75fb5a83d7caa'
    tags = ['infoleak:submission="manual"']
    name = 'submitted queue'
    description = 'first test queue, all submitted items'
    #queue_uuid = ''

    #res = create_ail_instance(ail_uuid, url, api_key=api_key, description=description)

    #res = create_sync_queue(name, tags=tags, description=description, max_size=100)
    #res = delete_sync_queue(queue_uuid)

    #res = register_ail_to_sync_queue(ail_uuid, queue_uuid)
    res = change_pull_push_state(ail_uuid, push=True, pull=True)

    print(res)