#!/usr/bin/python3
import os
import subprocess
import re
import sys
import time
import tempfile
import shutil
import gettext
import locale
from astra_update_service import ServiceLogger, LogSaver

translation_dir = "/usr/share/locale"
current_locale, encoding = locale.getlocale()

lang = gettext.translation('astra-update-service', localedir=translation_dir, languages=["ru"], fallback=True)
lang.install()
_ = lang.gettext

if current_locale != "ru_RU":
    _ = lambda s: s # if locale is not russian return original string

log_dir = "/var/log/astra-update-service"
log_file="/var/log/astra-update-service/update.log"
upgradable_stamp="/var/cache/astra-update-service/upgradable"
snapshot_stamp="/var/cache/astra-update-service/snapshot"
success_stamp="/var/cache/astra-update-service/successful"
failure_stamp="/var/cache/astra-update-service/unsuccessful"
main_app_dir="/var/cache/astra-update-service"
scripts_dir="/usr/share/astra-update-service/scripts"
sources_dir="/var/cache/astra-update-service/sources"
sorted_scripts_dir="/var/cache/astra-update-service/scripts"
service_conf_file="/etc/astra-update-service/astra-update-daemon.conf"
repos_policy_file="/var/cache/astra-update-service/repos_policy"

sys_upd_link="/system-update"
plymouth_running = False
conf_file = "/etc/astra-update-service/temp.ini"
boot_snap_path = "/var/tmp/boot_snap.tar.bz2"
xfs_snapshot_path = "/var/tmp/xfssnap-astra-update-service.bin"


logger = ServiceLogger(
    "astra-update-service",
    service_conf_file,
    log_file
)

def usage():
    # Help screen
    print(_("Please, don't operate this service manually."))
#    print('\t' + os.path.basename(__file__) + ' [command]')
#    print('Available commands:')
#    print('\t update \t Check updates')
#    print('\t download \t Download packages')
#    print('\t upgrade \t Upgrade packages')
#    print('\t clean \t\t Clean package archive')
#    print('\t check \t\t Check broken packages')
    return -1

def plymouth_check():
    global plymouth_running
    plymouth_run = subprocess.run(['plymouth', '--ping'])
    if plymouth_run.returncode == 0:
        plymouth_running = True
        logger.info('plymouth is running')
    else:
        logger.info('plymouth is not running')
        plymouth_running = False
                              
def display_message(text):
    if plymouth_running:
        #Print message on console or Plymouth
        subprocess.run(["plymouth", "message", "--text", text])

        #if plymouth_running:
        #    subprocess.run(["plymouth", "display-message", "--text", text])
    else:
        print(text)

def plymouth_update_mode ():
    #print("Setting system-upgrade mode")
    logger.debug("plymouth_update_mode")
    subprocess.run(["plymouth", "change-mode", "--updates"])

def plymouth_update_status(status):
    #print(percent)
    logger.debug("plymouth_update_status = " + str(status))
    subprocess.run(["plymouth", "update", "--status", status])

def plymouth_boot_mode ():
    #print("Setting boot-up mode")
    logger.debug("plymouth_boot_mode")
    subprocess.run(["plymouth", "change-mode", "--boot-up"])

def check_dirs ():
    # Check log dir, if dir not found, create it
    logger.debug("check_dirs")
    if os.path.exists(log_dir) == False:
        logger.debug("log_dir does not exist")
        os.mkdir(log_dir)
    if os.path.exists(main_app_dir) == False:
        logger.debug("main_app_dir does not exist")
        os.mkdir(main_app_dir)
    if os.path.exists(log_file) == False:
        logger.debug("log_file does not exist")
        open(log_file, 'a')

def clean_cache():
    logger.info("Clean apt caches")
    cleanproc = subprocess.run(['apt-get', 'clean'])
    if cleanproc.returncode != 0:
        logger.debug("cleanproc.returncode != 0, return 1")
        return 1
    cleanproc = subprocess.run(['apt-get', '-qq', '-y', 'autoclean'])
    if cleanproc.returncode != 0:
        logger.debug("cleanproc.returncode != 0, return 2")
        return 2
    return 0

def check_battery_charge():
    # Check battery charge. If charge < 50% - disable update.
    logger.debug("check_battery_charge")

    if os.path.exists("/sys/class/power_supply/BAT0"):
        logger.debug("os.path.exists(\"/sys/class/power_supply/BAT0\")")
        with open("/sys/class/power_supply/BAT0/capacity") as f:
            if int(f.read()) > 50:
                logger.info("power supply: battery capacity more than 50 percents")
                return 0
            else:
                logger.error("Battery charge less then 50%")
                return 1
    else:
        logger.warning("no information about battery")
        return 0

def check_power_supply() :
    logger.debug("check_power_supply")
    if os.path.exists("/sys/class/power_supply/AC"):
        logger.debug("os.path.exists(\"/sys/class/power_supply/AC\")")
        with open("/sys/class/power_supply/AC/online") as f:
            if int(f.read(1)) == 1:
                logger.info("power supply: AC")
                return 0
            else:
                logger.debug("check_power_supply: int(f.read(1)) != 1")
                return check_battery_charge()
    else:
        logger.warning("no information about ac")
        return 0


def update_symlink(create):
    # create or remove /system-update symlink
    logger.debug("update_symlink: create = " + str(create))
    if create == True:
        if not os.path.exists(sys_upd_link):
            logger.debug("not os.path.exists(sys_upd_link)")
            os.symlink("/var/cache/apt/archives", sys_upd_link)
        else:
            logger.warning("the symlink already exists and will be not created again")
                
    else:
        if os.path.exists(sys_upd_link):
            logger.debug("os.path.exists(sys_upd_link)")
            os.remove(sys_upd_link)
        else:
            logger.warning("the symlink not exists and will be not removed therefore")

def fix_dpkg():
    #fixing broken packages
    logger.info("Try to fix broken packages...")
    with tempfile.TemporaryFile(mode='w+') as temp_file:
        fixproc = subprocess.run(["dpkg", "--configure", "-a"], capture_output=True, text=True)
        res = fixproc.returncode
        if res != 0:
            logger.error("Can't fix broken packages")
            temp_file.seek(0)
            logger.error(temp_file.read())
            return 1
        else:
            logger.info("Broken packages fixed")
            return 0

def check_pkgmngr(repospolicy):
    # check packet manager on errors
    if os.path.islink(sys_upd_link):
        logger.debug("os.path.islink(sys_upd_link)")
        with tempfile.TemporaryFile(mode='w+') as temp_file:
            res = 0
            checkproc = subprocess.run(["dpkg", "--audit"], capture_output=True, text=True)
            res += checkproc.returncode
            temp_file.write(str(checkproc.stdout))
            checkproc = subprocess.run(["apt-get", "check"], capture_output=True, text=True)
            res += checkproc.returncode
            temp_file.write(str(checkproc.stdout))
            command = ["apt-get", "update"] + get_args_for_repos(repospolicy)
            checkproc = subprocess.run(command, capture_output=True, text=True)
            temp_file.write(str(checkproc.stdout))      
            if res != 0:
                logger.warning("Packet management system has problems")
                temp_file.seek(0)
                logger.warning(temp_file.read())                
                return fix_dpkg()
            logger.info("Package management is OK")
            return 0
    else:
        logger.warning(sys_upd_link + " checking failed")
    return 2

def checks_before_update(repospolicy):
    # check if the system is ready for update
    logger.info("Checking if the system is ready for update")
    if os.path.isfile(upgradable_stamp):
        if (check_power_supply() or check_pkgmngr(repospolicy)):
            update_symlink(False)
            logger.error("Errors detected during the system check, disabling update")
            return 1
        else:
            return 0
    else:
        logger.info("Upgradable stamp is missing or corrupted, an update will be cancelled")
        update_symlink(False)
        return 2

def get_args_for_repos(repospolicy):
    logger.debug("get_args_for_repos: repospolicy = " + str(repospolicy))
    if repospolicy == 0:
        return(["-o", "Dir::Etc::sourceparts=''"])
    if repospolicy == 1:
        return([])
    if repospolicy == 2:
        return(["-o", "Dir::Etc::sourceparts=" + sources_dir, "-o", "Dir::Etc::sourcelist=''"])
    if repospolicy == 3:
        return(["-o", "Dir::Etc::sourceparts=" + sources_dir])

def get_update_packages_list(repospolicy):
    # update packages list
    clearlist = []

    logger.debug("args = " + str(get_args_for_repos(repospolicy)))

    command = ["apt-get", "update"] + get_args_for_repos(repospolicy)
    subprocess.run(command)
    command = ["apt", "list", "--upgradable"] + get_args_for_repos(repospolicy)
    aptproc = subprocess.run(command, capture_output=True, text=True)
    #aptproc = subprocess.run(["apt", "list", "--upgradable"], capture_output=True, text=True)
    if aptproc.returncode == 0:
        logger.debug("aptproc.returncode == 0")
        list = aptproc.stdout.split("\n")
        for package in list[1:]:
            if (not package.isspace()) and (len(package) != 0) & (package.find("/") > 0): #TODO test &
                clearlist.append(package[0:package.find('/')])
    return clearlist

def check_updates(repospolicy):
    logger.debug("check_updates: repospolicy = " + str(repospolicy))

    command = ["apt-get", "update"] + get_args_for_repos(repospolicy)
    subprocess.run(command)
    command = ["apt", "list", "--upgradable"] + get_args_for_repos(repospolicy)
    aptproc = subprocess.run(command, capture_output=True, text=True)
    #aptproc = subprocess.run(["apt", "list", "--upgradable"], capture_output=True, text=True)
    if aptproc.returncode == 0:
        logger.debug("aptproc.returncode == 0")
        list = aptproc.stdout.split("\n")
        for package in list[1:]:
            if (not package.isspace()) and (len(package) != 0) & (package.find("/") > 0):
                return 1
    return 0
    
#def get_new_astra_version():
#    aptproc = subprocess.run(["apt-cache", "policy", "astra-version"], capture_output=True, text=True)
#    if aptproc.returncode == 0:
#        list = aptproc.stdout.split("\n")
#        if len(list) < 3:
#            return 1
#        versions = []
#        for package in list[1:3]:
#            if (not package.isspace()) and (len(package) != 0) :
#                split = package.split(":")
#                if len(split) >= 2:
#                    version = split[1].strip()
#                    match = re.search('\+v', version)
#                    if match:
#                        versions.append(version[match.end():])
#        if len(versions) == 2:
#            print(versions)
# #            if (versions[0] == versions[1]):
# #                return 2
# #            else:
#            pos = versions[1].rfind('.')
#            print(pos)
#            print(versions[1])
#
#            versions[1] = versions[1].replace('.', '')
#            print(versions[1])
#            print(int(versions[1][:pos-2]))
#            return int(versions[1][:pos-2])
#
#    return 3

def download_packages(repospolicy):
    # download  packages
    logger.debug("download_packages: repospolicy = " + str(repospolicy))
    packages_list = get_update_packages_list(repospolicy)
    if len(packages_list) == 0:
        logger.info("No packages to upgrade")
        file = open(conf_file, "w")
        file.write("2")
        file.close()
        return 2
    else:        
        logger.info("Start downloading updates")
        command = ["apt-get", "--download-only", "-y", "dist-upgrade"] + get_args_for_repos(repospolicy)
        aptproc = subprocess.Popen(command, stdout=subprocess.PIPE, text=True)
        while True:
            output = aptproc.stdout.readline()
            if output == '' and aptproc.poll() is not None:
                break
            if not output.strip():
                continue
            if output:
                match = re.search(r'\w\w\w:[0-9]', output)
                if match:
                    split = output.split(' ') 
                    if len(split) >= 7:
                        logger.info(str("Get: " + split[4] + "_" + split[6]))
                    else:
                        logger.warning("Strange info, full print:")
                        logger.warning(output)
        if aptproc.returncode == 0:
            logger.info("Downloading complete")
            file = open(conf_file, "w")
            file.write("0")
            file.close()
            return 0

        else:
            logger.error("Error while downloading packages:")
            for err_pack in packages_list:
                logger.error(err_pack)
            file = open(conf_file, "w")
            file.write("1")
            file.close()
            return 1
    
def get_download_result():
    logger.debug("get_download_result")
    with open(conf_file, "r") as f:
        for line in f:
            res = int(line)
    return res

def show_error_message(msg):
    logger.debug("show_error_message: " + msg)
    plymouth_update_status(msg)
    for i in range(5):
        msg = _("Reboot in ") + str(5 - i) + _(" s")
        plymouth_update_status(msg)
        time.sleep(1)

def get_fs_type():
    logger.debug("get_fs_type")

    global fs_type
    fs_type = -1
    global free_space
    free_space = -1
#    global mount_point
    global disk_name
    global disk_ok
    disk_ok = False

    dfproc = subprocess.run(["df", "-T", "-h"], capture_output=True, text=True)
    if dfproc.returncode == 0:
        logger.debug("dfproc.returncode == 0")
        list = dfproc.stdout.split("\n")
        for device in list[1:]:
            if (not device.isspace()) and (len(device) != 0):
                if re.match(".*\s\/$", device):
                    split = device.split()
                    if len(split) == 7:
                        disk_ok = True
                    else:
                        disk_ok = False
                        return
                    disk_name = split[0]
                    free_space = split[4]
#                    mount_point = split[6]
                    if split[1] == 'xfs':
                        logger.info("File system type: XFS")
                        fs_type = 0
                    elif '/dev/mapper/' in split[0]:
                        logger.info("File system type: LMV")
                        fs_type = 1
                    else:
                        logger.debug("File system type: other")
                        fs_type = 2
    else:
        logger.debug("dfproc.returncode != 0")
        disk_ok = False

def create_xfs_snapshot():
    plymouth_update_status(_("Creating snapshot"))
    logger.info("XFS snapshot creating")
    snap_proc = subprocess.Popen(["xfsdump",  "-l", "0", "-L", "session_astra", "-M",
                                    "tape_astra", "-f", xfs_snapshot_path, "/"],
                                    stdout=subprocess.PIPE, text=True)
#    subprocess.run(["xfsdump",  "-l", "0", "-L", "session_astra", "-M", "tape_astra", "-f", xfs_snapshot_path, "/"],
#                   stdout=subprocess.PIPE, text=True)
    #TODO add stdout processing
    while True:
        if snap_proc.poll() is None:
            plymouth_update_status("Snapshot is being created, please wait")
            logger.info("XFS snapshot creating is in progress")
            time.sleep(5)
        else:
            logger.info("XFS snapshot created: returncode = " + str(snap_proc.poll()))
            break

def get_diskname_for_lvm():
    m_disk_name = disk_name.replace("--", "|")
    m_disk_name = m_disk_name.replace("-", "/")
    m_disk_name = m_disk_name.replace("|", "-")
    m_disk_name = m_disk_name.replace("mapper/", "")
    logger.debug("get_diskname_for_lvm: m_disk_name = " + m_disk_name)
    return m_disk_name

def get_vg_section():
    match = re.search("^\/dev\/(?P<vg>.*)\/.+$", get_diskname_for_lvm())
    if match:
        m_vg = match.group('vg')
    else:
        m_vg = ""

    logger.debug("get_vg_section: m_vg = " + m_vg)
    return m_vg

def create_boot_snap() -> int:
    """
    Archivate /boot content.
    Return code:\n
    0 - success\n
    1 - archive created partially\n
    2 - fatal error creating archive
    """

    if (os.path.exists(boot_snap_path)):
        logger.error(f"create_boot_snap: {boot_snap_path} already exists, aborting")
        print("already exists")
        return 2

    boot_dir = "/boot"

    # List of all elems in /boot
    boot_contents = os.listdir(boot_dir)
    boot_contents_full_paths = [os.path.join(boot_dir, entry) for entry in boot_contents]

    tar_cmd = ["tar", "--xattrs", "--acls", "-czpf", boot_snap_path, "--exclude=/boot/lost+found"] + boot_contents_full_paths

    return subprocess.run(tar_cmd).returncode


def restore_snapshot():
    exec_current_stage_scripts(3, 1)
    get_fs_type()
    logger.debug("restore_snapshot: fs_type = " + str(fs_type))
    if (fs_type == 0):
        restore_xfs_snapshot()
    elif (fs_type == 1):
        restore_lvm_snapshot()
    else:
        plymouth_update_status(_("Snapshots are supported for LVM and XFS systems only"))
        logger.info("Snapshots are not supported for this system type")
    exec_current_stage_scripts(3, 0)

def restore_xfs_snapshot():
    log_saver = LogSaver()
    plymouth_update_status(_("Restoring snapshot"))
    logger.info("XFS snapshot restoring: " + xfs_snapshot_path)
    subprocess.run(["xfsrestore", "-f", xfs_snapshot_path, "/"])
    log_saver.restore()

def restore_lvm_snapshot():
    log_saver = LogSaver()
    plymouth_update_status(_("Restoring snapshot"))
    logger.info("LVM snapshot restoring")
    snapproc = subprocess.run(["/usr/bin/astra-update-modern", "-br"], capture_output=True, text=True)
    logger.debug(f"restore_lvm_snapshot: snapproc.returncode = {snapproc.returncode}")
    if snapproc.returncode != 0:
        logger.error("astra-update-modern -br return error: " + str(snapproc.stderr) + "\n" + str(snapproc.stdout))
        return
    elif os.path.exists(boot_snap_path):
        clear_directory("/boot")
        tarproc = subprocess.run(["tar", "--xattrs", "--xattrs-include=security.{PDPL,AUDIT,DEF_AUDIT}", "--acls", "-xzf", boot_snap_path, "-C", "/"])
        logger.debug(f"restore_lvm_snapshot: tarproc.returncode = {tarproc.returncode}")
        if tarproc.returncode != 0:
            logger.error(f"restore_lvm_snapshot: error retrieving boot snapshot. tarproc.returncode = {tarproc.returncode}")
        else:
            os.remove(boot_snap_path)

    log_saver.restore()


def clear_directory(path):
    """
    Removes all files and dirs in provided directory but
    does not remove the provided directory.
    """
    for filename in os.listdir(path):
        try:
            file_path = os.path.join(path, filename)
            if os.path.isfile(file_path) or os.path.islink(file_path):
                os.remove(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        except Exception as e:
            logger.error(f'clear_directory: Error deleting {file_path}')
            logger.error(f'clear_directory: Exception type - {str(type(e).__name__)}')
            logger.error(f'clear_directory: Message: {str(e)}')


def remove_snapshot():
    get_fs_type()
    logger.debug("remove_snapshot: fs_type = " + str(fs_type))
    if (fs_type == 0):
        remove_xfs_snapshot()
    elif (fs_type == 1):
        remove_lvm_snapshot()
    else:
        plymouth_update_status(_("Snapshots are supported for LVM and XFS systems only"))
        logger.info("Snapshots are not supported for this system type")

def remove_xfs_snapshot():
    plymouth_update_status(_("Removing snapshot"))
    logger.info("XFS snapshot deleting: " + xfs_snapshot_path)
    if (os.path.exists(xfs_snapshot_path)):
        os.remove(xfs_snapshot_path)
        logger.info("Snapshot has been deleted")
    else:
        logger.warning("Snapshot file hadn't been found")

def remove_lvm_snapshot():
    plymouth_update_status(_("Removing snapshot"))
    logger.info("LVM snapshot deleting")
    snapproc = subprocess.run(["/usr/bin/astra-update-modern", "-bd"], capture_output=True, text=True)

    if snapproc.returncode != 0:
        logger.error("error removing lvm snapshot: " + str(snapproc.stderr) + "\n" + str(snapproc.stdout))

    if os.path.exists(boot_snap_path):
        os.remove(boot_snap_path)
    else:
        logger.warning(boot_snap_path + " hadn't been found")

def get_repos_policy():
    logger.debug("get_repos_policy")
    repospolicy = 0
    if os.path.isfile(repos_policy_file):
        logger.debug("os.path.isfile(repos_policy_file)")
        with open(repos_policy_file, 'r') as f:
            data = f.read()
            repospolicy = int(data)
    return repospolicy

def reboot_with_error(errmsg: str):
    show_error_message(errmsg)
    if os.path.exists(main_app_dir) == False:
        os.mkdir(main_app_dir)
    new_stamp = open(failure_stamp, 'w')
    new_stamp.write("0")
    new_stamp.close()
    logger.warning("Rebooting system due errors before update")
    subprocess.run(["systemctl", "reboot"])

def upgrade():
    logger.debug("upgrade")
    repospolicy = get_repos_policy()
    plymouth_check()
    exec_current_stage_scripts(2, 1)
    if (checks_before_update(repospolicy) != 0):
        reboot_with_error(_("Error of update checks"))
        return
    plymouth_update_mode()
    plymouth_update_status(_("System update in progress. Do not turn off the computer!"))

    need_to_create_lvm_snapshot = False
    if os.path.exists(snapshot_stamp):
        get_fs_type()
        logger.debug("create_snapshot: fs_type = " + str(fs_type))
        #xfs
        if (fs_type == 0):
            create_xfs_snapshot()
        #lvm
        elif (fs_type == 1):
            returncode = create_boot_snap()

            if (returncode == 0):
                need_to_create_lvm_snapshot = True
            elif (returncode == 1):
                os.remove(boot_snap_path)

            if (returncode != 0):
                logger.error(f"can not create boot snapshot. tar returncode = {returncode}")
        else:
            plymouth_update_status(_("Snapshots are supported for LVM and XFS systems only"))
            logger.info("Snapshots are not supported for this system type")

    space_policy = get_space_checking_policy()
    start_msg = _(("System update in progress. Do not turn off the computer: "))
    logger.info("Symlink removing...")
    update_symlink(False)
    logger.info("Start system update")
    command = ["/usr/bin/astra-update-modern", "-A", "-N"]
    logger.debug("need_to_create_lvm_snapshot = " + str(need_to_create_lvm_snapshot))
    if need_to_create_lvm_snapshot:
        command.append("-bc")
    #if space_policy == 0:
        #command = ["astra-update", "-A", "-T", "-r", "-N", "-o", "APT::Status-Fd=2"] + get_args_for_repos(repospolicy)
        #aptget_proc = subprocess.Popen(["astra-update", "-A", "-T", "-r", "-N", "-o", "APT::Status-Fd=2"], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
    #else:
    #    command = ["astra-update", "-A", "-T", "-r", "-N", "-S", "-o", "APT::Status-Fd=2"] + get_args_for_repos(repospolicy)
    logger.debug("space_policy = " + str(space_policy) + " repospolicy = " + str(repospolicy))
    if space_policy != 0:
        command += ["-S"]
    if repospolicy == 1:
        command += ["-r", "-T"]
    command += ["-o", "APT::Status-Fd=2"]
    if repospolicy != 1:
        command += get_args_for_repos(repospolicy)
    logger.debug("full command: " + str(command))
    aptget_proc = subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
#        aptget_proc = subprocess.Popen(["astra-update", "-A", "-T", "-r", "-N", "-S", "-o", "APT::Status-Fd=2"], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
#    aptget_proc = subprocess.Popen(["apt-get", "-f", "-y", "-o", "Apt::Status-fd=2", "-o",
#                                    "Dpkg::Options::=--force-confdef", "-o", "Dpkg::Options::=--force-confold",
#                                    "dist-upgrade"], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
    while True:
        errdata = aptget_proc.stderr.readline()
        if errdata == '' and aptget_proc.poll() is not None:
            break
        if not errdata.strip():
            continue
        if errdata:
            split = errdata.split(':')
            if len(split) >= 4:
                logger.info(str(split[2] + "% " + split[3]))
                int_percent = split[2].split('.')
                try:
                    percent = float(split[2])
                except ValueError:
                    percent = 0
                plymouth_update_status(start_msg + str(round(percent, 2)) + "%" + "|" + int_percent[0])
            else:
                logger.warning("Strange info, full print:")
                logger.warning(errdata)
    if aptget_proc.returncode!=0:
        errmsg = _("Error during upgrade. Contact your system administrator")
        show_error_message(errmsg)
        if os.path.exists(main_app_dir) == False:
            os.mkdir(main_app_dir)
        new_stamp = open(failure_stamp, 'w')
        new_stamp.write("1")
        new_stamp.close()
        logger.error("Error during upgrade. Contact your system administrator")
        logger.error(str(aptget_proc.stderr))
        logger.warning("Rebooting system due errors")
        plymouth_boot_mode()
        subprocess.run(["systemctl", "reboot"])
    else:
        logger.info("Autoremove unused packages")
        plymouth_update_status(_("Clearing Cache"))
        os.remove(upgradable_stamp)
        clean_cache()
        if os.path.exists(main_app_dir) == False:
            os.mkdir(main_app_dir)
        new_stamp = open(success_stamp, 'w')
        new_stamp.close()
        exec_current_stage_scripts(2, 0)
        logger.info("Rebooting system")
        plymouth_boot_mode()
        subprocess.run(["systemctl", "reboot"])

def get_space_checking_policy():
    logger.debug("get_space_checking_policy")
    with open(service_conf_file, "r") as f:
        for line in f:
            if line.startswith("Free_space_policy"):
                if len(line.split("=")) > 1 :
                    print(line.split("=")[1].strip())
                    return int(line.split("=")[1].strip())
    return 0

def get_required_space_in_stdout(repospolicy):
    logger.debug("get_required_space_in_stdout: " + str(repospolicy))

    apt_args = get_args_for_repos(repospolicy)
    command = ["/usr/sbin/apt-cache-calculator", "--apt_args"] + apt_args
    result = subprocess.run(command, capture_output=True, text=True)

    if (result.returncode == 1):
        logger.debug("get_required_space_in_stdout: len(packages_list) == 0")
        return 1

    logger.debug("len(packages_list) != 0")
    size_to_install_Mb = result.stdout.split('\n')[0]
    print(size_to_install_Mb)
    return 0

def check_arg():
    logger.debug("check_arg: sys.argv = " + ' '.join(sys.argv))
    if len(sys.argv) - 1 == 0:
        return usage()
    else:
        if sys.argv[1] == "download":
            if len(sys.argv) == 3:
                return download_packages(int(sys.argv[2]))
            else:
                return download_packages(0)
        if sys.argv[1] == "upgrade":
            return upgrade()
        if sys.argv[1] == "clean":
            return clean_cache() #not sure            
#        if sys.argv[1] == "check":
#            return checks_before_update() //add repospolicy if you want to use it
        if sys.argv[1] == "update":
            if len(sys.argv) == 3:
                return check_updates(int(sys.argv[2]))
            else:
                return check_updates()
        if sys.argv[1] == "downresult":
            return get_download_result()
        if sys.argv[1] == "rm_snapshot":
            return remove_snapshot()
        if sys.argv[1] == "rollback":
            return restore_snapshot()
 #       if sys.argv[1] == "get_version":
 #           return get_new_astra_version()
        if sys.argv[1] == "get_required_space":
            if len(sys.argv) == 3:
                return get_required_space_in_stdout(int(sys.argv[2]))
            else:
                return get_required_space_in_stdout(0)
        if sys.argv[1] == "exec_scripts":
            if len(sys.argv) < 4:
                return usage()
            else:
                try:
                    prev_stage = ""
                    stage = int(sys.argv[2])
                    pre_stage = int(sys.argv[3])
                    if (stage >= 0 and stage < 4) and (pre_stage == 0 or pre_stage == 1):
                        if (len(sys.argv) == 5):
                            prev_stage = int(sys.argv[4])
                            return(exec_current_stage_scripts(stage, pre_stage, prev_stage))
                        else:
                            return(exec_current_stage_scripts(stage, pre_stage))
                    else:
                        return usage()
                except ValueError:
                    return usage()
        return usage()

def create_script_dirs(): #do we need to clean existing folders?
    if os.path.exists(sorted_scripts_dir) == False:
        logger.debug("os.path.exists(sorted_scripts_dir) == False")
        os.mkdir(sorted_scripts_dir)
    dirname = sorted_scripts_dir + "/pre_ready"
    if os.path.exists(dirname) == False:
        logger.debug(f'os.path.exists({dirname}) == False')
        os.mkdir(dirname)
    dirname = sorted_scripts_dir + "/post_ready"
    if os.path.exists(dirname) == False:
        logger.debug(f'os.path.exists({dirname}) == False')
        os.mkdir(dirname)
    dirname = sorted_scripts_dir + "/pre_activated"
    if os.path.exists(dirname) == False:
        logger.debug(f'os.path.exists({dirname}) == False')
        os.mkdir(dirname)
    dirname = sorted_scripts_dir + "/post_activated"
    if os.path.exists(dirname) == False:
        logger.debug(f'os.path.exists({dirname}) == False')
        os.mkdir(dirname)
    dirname = sorted_scripts_dir + "/pre_upgrade"
    if os.path.exists(dirname) == False:
        logger.debug(f'os.path.exists({dirname}) == False')
        os.mkdir(dirname)
    dirname = sorted_scripts_dir + "/post_upgrade"
    if os.path.exists(dirname) == False:
        logger.debug(f'os.path.exists({dirname}) == False')
        os.mkdir(dirname)
    dirname = sorted_scripts_dir + "/pre_rollback"
    if os.path.exists(dirname) == False:
        logger.debug(f'os.path.exists({dirname}) == False')
        os.mkdir(dirname)
    dirname = sorted_scripts_dir + "/post_rollback"
    if os.path.exists(dirname) == False:
        logger.debug(f'os.path.exists({dirname}) == False')
        os.mkdir(dirname)

def get_script_info(script_path):
    logger.debug(f'get_script_info: scripts_path = {script_path}')
    checkproc = subprocess.run([script_path, "config"], capture_output=True, text=True)
    pre_state = -1
    cout_list = checkproc.stdout.split("\n")
    for str in cout_list:
        if pre_state < 0:
            if str.find("substage: pre") == 0:
                pre_state = 1
            elif str.find("substage: post") == 0:
                pre_state = 0
    for str in cout_list:
        if str.find("stage") == 0:
            if str.find("ready") > 0:
                copy_script_to_folder(script_path, 0, pre_state)
            if str.find("activated") > 0:
                copy_script_to_folder(script_path, 1, pre_state)
            if str.find("upgrade") > 0:
                copy_script_to_folder(script_path, 2, pre_state)
            if str.find("rollback") > 0:
                copy_script_to_folder(script_path, 3, pre_state)

def generate_folder_by_stage(stage, pre_state):
    logger.debug(f'generate_folder_by_stage: stage = {str(stage)}, pre_state = {str(pre_state)}')
    folder_name = ""
    if stage == 0:
        folder_name = "ready"
    elif stage == 1:
        folder_name = "activated"
    elif stage == 2:
        folder_name = "upgrade"
    elif stage == 3:
        folder_name = "rollback"
    if pre_state == 0:
        folder_name = "/post_" + folder_name
    else:
        folder_name = "/pre_" + folder_name
    return sorted_scripts_dir + folder_name

def copy_script_to_folder(script_name, stage, pre_state):
    logger.debug(f'copy_script_to_folder: script_name = {script_name}, stage = {str(stage)}, pre_state = {str(pre_state)}')
    abs_folder_name = generate_folder_by_stage(stage, pre_state)
    shutil.copy(script_name, abs_folder_name)

def load_scripts():
    logger.debug("load_scripts")
    for root, dirs, files in os.walk(scripts_dir):
        if len(files) > 0:
            create_script_dirs()
        for script in files:
            get_script_info(os.path.join(root, script))

def exec_current_stage_scripts(stage, pre_state, prev_stage = None):
    logger.debug(f'copy_script_to_folder: stage = {str(stage)}, pre_state = {str(pre_state)}, prev_stage = {str(prev_stage)}')
    use_prev_name = False
    if (prev_stage is not None):
        use_prev_name = True

    stage_name = ""
    if stage == 0:
        stage_name = "ready"
    elif stage == 1:
        stage_name = "activated"
    elif stage == 2:
        stage_name = "upgrade"
    elif stage == 3:
        stage_name = "rollback"

    if pre_state == 0:
        substage_name = "post"
    else:
        substage_name = "pre"

    prev_stage_name = ""
    if prev_stage == 0:
        prev_stage_name = "ready"
    elif prev_stage == 1:
        prev_stage_name = "activated"
    elif prev_stage == 2:
        prev_stage_name = "upgrade"
    elif prev_stage == 3:
        prev_stage_name = "rollback"

    logger_text = "Scripts execution: stage = " + stage_name + " sub_stage = " + substage_name
    if use_prev_name:
        logger_text += " prev_stage = " + prev_stage_name
    logger.info(logger_text)
    current_folder = generate_folder_by_stage(stage, pre_state)
    summ_errors = 0
    for root, dirs, files in os.walk(current_folder):
        if len(files) == 0:
            logger.warning("No scripts were found")
            return
        for script in files:
            curr_error = 0
            logger.info("Current script: " + script)
            script_path = current_folder + "/" + script
            if use_prev_name:
                logger.debug("use_prev_name = true")
                scriptproc = subprocess.run([script_path, "prev", prev_stage_name], capture_output=True, text=True)
            else:
                logger.debug("use_prev_name = false")
                scriptproc = subprocess.run([script_path], capture_output=True, text=True)
            cout_list = scriptproc.stdout.split("\n")
            for out_str in cout_list:
                if out_str.find("message:") >= 0:
                    logger.info(out_str)
                if out_str.find("warning:") >= 0:
                    logger.warning(out_str)
                if out_str.find("error:") >= 0:
                    logger.error(out_str)
                    curr_error = 1
            if scriptproc.returncode == 0:
                logger.info("Script return code = 0")
            else:
                curr_error = 1
                logger.info("Script return code = " + str(scriptproc.returncode))
            if curr_error == 1:
                summ_errors = summ_errors + 1
    return summ_errors

check_dirs()
load_scripts()
exit(check_arg())
