AIL-framework/bin/lib/ail_updates.py

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

import datetime
import os
import sys

sys.path.append(os.environ['AIL_BIN'])
##################################
# Import Project packages
##################################
from lib.ConfigLoader import ConfigLoader
from lib import Users

config_loader = ConfigLoader()
#r_serv_db = config_loader.get_redis_conn("Kvrocks_DB")
r_serv = config_loader.get_redis_conn("ARDB_DB")
config_loader = None

BACKGROUND_UPDATES = {
    'v1.5':{
            'nb_updates': 5,
            'message': 'Tags and Screenshots'
           },
    'v2.4':{
            'nb_updates': 1,
            'message': ' Domains Tags and Correlations'
           },
    'v2.6':{
            'nb_updates': 1,
            'message': 'Domains Tags and Correlations'
           },
    'v2.7':{
            'nb_updates': 1,
            'message': 'Domains Tags'
           },
    'v3.4':{
            'nb_updates': 1,
            'message': 'Domains Languages'
           },
    'v3.7':{
            'nb_updates': 1,
            'message': 'Trackers first_seen/last_seen'
           }
}

def get_ail_version():
    return r_serv.get('ail:version')

def get_ail_float_version():
    version = get_ail_version()
    if version:
        version = float(version[1:])
    else:
        version = 0
    return version

def get_ail_all_updates(date_separator='-'):
    dict_update = r_serv.hgetall('ail:update_date')
    if date_separator:
        for version in dict_update:
            u_date = dict_update[version]
            dict_update[version] = f'{u_date[0:4]}{date_separator}{u_date[4:6]}{date_separator}{u_date[6:8]}'
    return dict_update

def add_ail_update(version):
    #Add new AIL version
    r_serv.hset('ail:update_date', version, datetime.datetime.now().strftime("%Y%m%d"))
    #Set current ail version
    if float(version[1:]) > get_ail_float_version():
        r_serv.set('ail:version', version)

def check_version(version):
    if version[0] != 'v':
        return False
    try:
        int(version[1])
        int(version[-1])
        int(version[1:].replace('.', ''))
    except:
        return False
    if '..' in version:
        return False
    return True

#### UPDATE BACKGROUND ####

def exits_background_update_to_launch():
    return r_serv.scard('ail:update:to_update') != 0

def is_version_in_background_update(version):
    return r_serv.sismember('ail:update:to_update', version)

def get_all_background_updates_to_launch():
    return r_serv.smembers('ail:update:to_update')

def get_current_background_update():
    return r_serv.get('ail:update:update_in_progress')

def get_current_background_update_script():
    return r_serv.get('ail:update:current_background_script')

def get_current_background_update_script_path(version, script_name):
    return os.path.join(os.environ['AIL_HOME'], 'update', version, script_name)

def get_current_background_nb_update_completed():
    return r_serv.scard('ail:update:update_in_progress:completed')

def get_current_background_update_progress():
    progress = r_serv.get('ail:update:current_background_script_stat')
    if not progress:
        progress = 0
    return int(progress)

def get_background_update_error():
    return r_serv.get('ail:update:error')

def add_background_updates_to_launch(version):
    return r_serv.sadd('ail:update:to_update', version)

def start_background_update(version):
    r_serv.delete('ail:update:error')
    r_serv.set('ail:update:update_in_progress', version)

def set_current_background_update_script(script_name):
    r_serv.set('ail:update:current_background_script', script_name)
    r_serv.set('ail:update:current_background_script_stat', 0)

def set_current_background_update_progress(progress):
    r_serv.set('ail:update:current_background_script_stat', progress)

def set_background_update_error(error):
    r_serv.set('ail:update:error', error)

def end_background_update_script():
    r_serv.sadd('ail:update:update_in_progress:completed')

def end_background_update(version):
    r_serv.delete('ail:update:update_in_progress')
    r_serv.delete('ail:update:current_background_script')
    r_serv.delete('ail:update:current_background_script_stat')
    r_serv.delete('ail:update:update_in_progress:completed')
    r_serv.srem('ail:update:to_update', version)

def clear_background_update():
    r_serv.delete('ail:update:error')
    r_serv.delete('ail:update:update_in_progress')
    r_serv.delete('ail:update:current_background_script')
    r_serv.delete('ail:update:current_background_script_stat')
    r_serv.delete('ail:update:update_in_progress:completed')

def get_update_background_message(version):
    return BACKGROUND_UPDATES[version]['message']

# TODO: Detect error in subprocess
def get_update_background_metadata():
    dict_update = {}
    version = get_current_background_update()
    if version:
        dict_update['version'] = version
        dict_update['script'] = get_current_background_update_script()
        dict_update['script_progress'] = get_current_background_update_progress()
        dict_update['nb_update'] = BACKGROUND_UPDATES[dict_update['version']]['nb_updates']
        dict_update['nb_completed'] = get_current_background_nb_update_completed()
        dict_update['progress'] = int(dict_update['nb_completed']*100/dict_update['nb_update'])
        dict_update['error'] = get_background_update_error()
    return dict_update


##-- UPDATE BACKGROUND --##








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

if __name__ == '__main__':
    res = check_version('v3.1..1')
    print(res)