#!/usr/bin/python3

# date: summer 2025
# python: 3.12.3
# OS: Ubuntu 24.04

import argparse
import time
import os
import socket
import subprocess
import logging

#----------------------------------------------------------------------

CUBE_SSH_USER = 'root'                                  # ssh-login
CUBE_SSH_KEY = '~/.ssh/cube_ssh_keys/v_id_rsa_aqc_root' # ssh-key
CUBE_UPDATER_FILE_BEFORE = '/data/updater-user.conf'    # файл конфигурации updater ДО обновления 
CUBE_UPDATER_FILE_AFTER = '/data/updater-user.conf'     # файл конфигурации updater ПОСЛЕ обновления 
CUBE_UPDATE_LINKS_FILE = '/data/update.links'           # файл ссылок при обновлении по http (на cube)
CUBE_UPTIME_AFTER = '0:02'                              # cube uptime ПОСЛЕ обовления (минуты)

SHELL_UPDATE_LINKS_FILE = '/tmp/update.links'           # файл ссылок при обновлении по http (на компьютере)
SHELL_LOG_FILE = '/tmp/local_update_helper_log.txt'     # лог-файл работы скрипта
SHELL_UPDATE_DIR = './swu-packet-updates'               # директория запуска http-сервера с прошивками

SCRIPT_PAUSE_BEFORE_VERIFY = 15                         # пауза в работе скрипта перед проверкой корректного обновления (минуты)
PAUSE_BETWEEN_ATTEMPT = 1                               # пауза перед началом очередной проверки корректного обновления (минуты)
VERIFY_ATTEMPT_MAX_COUNT = 5                            # количество проверок обновления

#----------------------------------------------------------------------
# Всё в папке /tmp считается временным, папка очищается при перезагрузке!!!
# Этот скрипт запускается на linux-компьютере, где будет http-сервер. Сам файл скрипта сделать исполняемым.
#
# Скрипт на основании содержимого локальной директории SHELL_UPDATE_DIR формирует файл для обновления по http 
# и сохраняет его по пути SHELL_UPDATE_LINKS_FILE. Затем этот файл копируется на каждый cube в 
# CUBE_UPDATE_LINKS_FILE и в файле CUBE_UPDATER_FILE_BEFORE заполняется поле Links_File. 
# Cube дается команда о перезагрузке.
# После паузы SCRIPT_PAUSE_BEFORE_VERIFY начинаем проверку факта обновления. Успешным обновлением считаем:
# 1) работу cube больше, чем CUBE_UPTIME И 2) не совпадающий с начальным, md5-хеш файла с 
# с установленными пакетами (файл /data/sw-versions).
# После успешного обновления очищаем поле Links_File в файле CUBE_UPDATER_FILE_AFTER
#
# Процесс автоматизации обновления
# 1. Запускаем http-сервер из директории с пакетами обновлений (из директории SHELL_UPDATE_DIR)
#
#    user@hostname:~/swu-packet-updates$ python3 -m http.server 8005
#
#    браузером заходим на страницу http://localhost:8005 или http://<your_ip>:8005 и убеждаемся, что сервер запущен
# 2. Запускаем данный скрипт - процесс обновления начнется. Например:
#
#    user@hostname:~$./local_update_helper.py --cube_ip 192.168.1.3 192.168.1.7 192.168.1.8
#
#
#
#
#------------ Helper --------------------------------------------------
class Helper:

    @staticmethod
    def isPing(ip: str) -> bool:
        try:
            result = subprocess.run(
                ['ping', '-c', '1', ip],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                timeout=3
            )
            if result.returncode == 0:
                return True
            else:
                return False
        except:
            return False

#----------------------------------------------------------------------

    @staticmethod
    def getOwnLanIP() -> str:
        fake_connection = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        fake_connection.settimeout(0)
        try:
            # организуем иммитацию соединения, адреса может и не существовать
            fake_connection.connect(('192.168.0.1', 1))
            ip = fake_connection.getsockname()[0]
        except:
            ip = '127.0.0.1'
        finally:
            fake_connection.close()
        return str(ip)

#----------------------------------------------------------------------

    @staticmethod
    def isValidIP(ip: str) -> bool:
        try:
            socket.inet_aton(ip)
            return True
        except:
            return False

#----------------------------------------------------------------------

    @staticmethod
    def isSSH(ip: str, ssh_user: str, ssh_key: str) -> bool:
        result = False
        try:
            response = subprocess.run(
                ['ssh', '-o', 'StrictHostKeyChecking=no', '-i', ssh_key, f'{ssh_user}@{ip}', 'exit'],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    timeout=5
                )
            if response.returncode == 0:
                result = True
        except:
            result = False
        return result

#----------------------------------------------------------------------
#------------ Log -----------------------------------------------------
class Log:


    def __init__(self, log_file: str):
        logging.basicConfig(filename=log_file,
                            format='%(asctime)s - %(levelname)s: %(message)s',
                            filemode='a',
                            encoding='utf-8',
                            level=logging.DEBUG)
        self.append('info', '=== START ==============================')

#----------------------------------------------------------------------

    def __del__(self):
        self.append('info', '=== STOP ===============================')

#----------------------------------------------------------------------

    def append(self, level: str, log_text: str, is_print_console: bool = False):
        severity = 'INFO   '
        match level:
            case 'debug':
                logging.debug(log_text)
                severity = 'DEBUG  '
            case 'warning':
                logging.warning(log_text)
                severity = 'WARNING'
            case 'error':
                logging.error(log_text)
                severity = 'ERROR  '
            case _:
                logging.info(log_text)

        if is_print_console:
            print(severity + ' - ' + log_text)

#----------------------------------------------------------------------

#------------ Shell ---------------------------------------------------
class Shell:
#----------------------------------------------------------------------

    def __init__(self, logFile: str, args):
        self.log = Log(logFile)
        self.args = args

#----------------------------------------------------------------------

    def isValidServer(self) -> bool:
        result = True

        if self.args.server_type[0] == 'ftp':
            self.args.server_type = self.args.server_type[0]
            self.log.append('error', 'FTP-сервер на текущий момент не поддерживается.', True)
            result = False

        if not os.path.isdir(self.args.server_dir):
            self.log.append('error', 'Директория сервера не существует.', True)
            result = False
        else:
            self.args.server_dir = os.path.abspath(self.args.server_dir)

        if not Helper.isPing(self.args.server_ip):
            self.log.append('error', 'IP-адрес сервера не пингуется.', True)
            result = False

        try:
            port = int(self.args.server_port)
            if port <= 1024 or port > 65535:
                result = False
        except:
            result = False

        return result

#----------------------------------------------------------------------

    def isCreateHTTPUpdateLinksFile(self, file: str) -> bool:
        """
        Создает файл с ссылками на файлы обновлений.
        """
        update_links_lines = []
        for dir_item in os.listdir(self.args.server_dir):
            if dir_item.endswith(".swu.bin"):
                unit_name = dir_item.split("_")[0]
                version = dir_item.split("_")[1].split(".swu.bin")[0]
                url = f"http://{self.args.server_ip}:{self.args.server_port}/{dir_item}"
                update_links_lines.append(f"{unit_name} {version} {url}")

        # Создаем временный файл update.links
        try:
            with open(file, "w") as file_handler:
                file_handler.write("\n".join(update_links_lines))
        except:
            file_handler.close()
            return False
        file_handler.close()
        return True

#----------------------------------------------------------------------
#------------ Cube ----------------------------------------------------
class Cube:


    def __init__(self, ip: str, login: str, key: str, updater_user_file: str):
        self.ip = ip
        self.ssh_user = login
        self.ssh_key = key
        self.updater_file = updater_user_file
        self.status = "find"                   # find/updating/updated/skip
        self.md5_hash_prev = self.getHash()
        self.md5_hash = None

#----------------------------------------------------------------------

    def runCommand(self, cube_command: str) -> str:
        ssh_connection =  f"ssh -o StrictHostKeyChecking=no -i {self.ssh_key} {self.ssh_user}@{self.ip}"
        try:
            response = subprocess.run(f"{ssh_connection} '{cube_command}'",
                                    stdout=subprocess.PIPE,
                                    stderr=subprocess.PIPE,
                                    shell=True,
                                    text=True
            )
            if response.returncode == 0:
                return response.stdout
            else:
                return response.stderr

        except:
            return ''

#----------------------------------------------------------------------

    def isReplaceLine(self, find_line_pattern: str, new_line) -> bool:
        try:
            run_command = f"scp -i {self.ssh_key} {self.ssh_user}@{self.ip}:{self.updater_file} /tmp/updater-user.back"
            response = subprocess.run(run_command,
                           stdout=subprocess.PIPE,
                           shell=True,
                           check=True)
            if response.returncode != 0:
                return False

            updater_file_handler = open('/tmp/updater-user.back', 'rt')
            text = updater_file_handler.read()
            text = text.replace(find_line_pattern, new_line)
            updater_file_handler.close()
            updater_file_handler = open('/tmp/updater-user.back', 'wt')
            updater_file_handler.write(text)
            updater_file_handler.close()

            run_command = f"scp -i {self.ssh_key} /tmp/updater-user.back {self.ssh_user}@{self.ip}:{self.updater_file}"
            response = subprocess.run(run_command,
                           stdout=subprocess.PIPE,
                           shell=True,
                           check=True)
            if response.returncode != 0:
                return False

            os.remove('/tmp/updater-user.back')
        except:
            return False

        return True
#----------------------------------------------------------------------

    def isPreparedUpdateHTTP(self, shell_update_links_file: str, cube_update_links_file: str) -> bool:
        try:
            # Копируем файл update.links на устройство
            run_command = f"scp -i {self.ssh_key} {shell_update_links_file} {self.ssh_user}@{self.ip}:{cube_update_links_file}"
            response = subprocess.run(run_command,
                           stdout=subprocess.PIPE,
                           shell=True,
                           check=True)
            if response.returncode != 0:
                return False

            # Удаленно меняем строку Links_File = "" в файле конфигурации updater
            if not self.isReplaceLine( 'Links_File = ""', f'Links_File = "{cube_update_links_file}"'):
                return False

        except:
            return False

        return True

#----------------------------------------------------------------------

    def getHash(self) -> str:
        result = self.runCommand('tail -n 1 /data/sw-versions')
        return result

#----------------------------------------------------------------------

    def reboot(self) -> None:
        self.runCommand("reboot")

#----------------------------------------------------------------------

    def uptime(self) -> str:
        result = self.runCommand('uptime')
        result = result.split(' ')
        result[4] = result[4].rstrip(',')
        return result[4]

#----------------------------------------------------------------------




#--- Проверяем вводные параметры ----------------------------------------

def main():
    parser = argparse.ArgumentParser(description='Скрипт обновления установленных програмных юнитов по локальной сети')
    parser.add_argument('--server_type', type=str, nargs=1, choices=['http', 'ftp'], default='http', help="http или ftp (умолч. http)")
    parser.add_argument('--server_ip', type=str, nargs=1, default=Helper.getOwnLanIP(), help='IP-адрес сервера обновлений (умолч. собственный ip-адрес)')
    parser.add_argument('--server_port', type=int, nargs=1, default='8005', help='Порт сервера обновлений (умолч. 8005, 1024 < server_port < 65535)')
    parser.add_argument('--server_dir', type=str, nargs=1, default='./swu-packet-updates', help='Локальная директория с пакетами обновлений на сервере обновлений (умолч. ./swu-packet-updates)')
    parser.add_argument('--cube_ip', type=str, nargs='+', help='IP-адрес обновляемого кубика (обязательный параметр)')

    shell = Shell(SHELL_LOG_FILE, parser.parse_args())

    if isinstance(shell.args.server_type, list):
        shell.args.server_type = shell.args.server_type[0]
    if isinstance(shell.args.server_ip, list):
        shell.args.server_ip = shell.args.server_ip[0]
    if isinstance(shell.args.server_port, list):
        shell.args.server_port = shell.args.server_port[0]



    shell.log.append('info', f'Лог-файл {SHELL_LOG_FILE}', True)
    if not shell.isValidServer():
        shell.log.append('error','Ошибка в параметрах сервера', True)
        return 1
    if shell.args.server_type == 'http':
        if not shell.isCreateHTTPUpdateLinksFile(SHELL_UPDATE_LINKS_FILE):
            shell.log.append('error', f'Ошибка при создании файла обновления по http {SHELL_UPDATE_LINKS_FILE}', True)
            return 1
        
#--- Создаем наши объекты-кубики в скрипте ------------------------------   
    
    addresses = list()
    cubes = list()
    continue_run = False
    

    for ip in shell.args.cube_ip:
        if ip not in addresses:
            if Helper.isPing(ip): 
                if Helper.isSSH(ip, CUBE_SSH_USER, CUBE_SSH_KEY):
                    addresses.append(ip)
                    cubes.append(Cube(ip, CUBE_SSH_USER, CUBE_SSH_KEY, CUBE_UPDATER_FILE_BEFORE))
                    shell.log.append('info', f'{ip} - в локальной сети обнаружен сube', True)
    
    shell.log.append('info', '========================================', True)
    shell.log.append('info', 'Параметры работы скрипта:', True)
    shell.log.append('info', f'Сервер обновений             - {shell.args.server_type}://{shell.args.server_ip}:{shell.args.server_port}', True)
    shell.log.append('info', f'Локальная директория сервера - {shell.args.server_dir}', True)
    shell.log.append('info', f'Всего cubes в работе         - {addresses}', True)
    shell.log.append('info', '========================================', True)

#--- Подготавливаем кубик к обновлению и ребутим его --------------------
    
    for cube in cubes:
        
        match shell.args.server_type:
            case 'http':
                if cube.isPreparedUpdateHTTP(SHELL_UPDATE_LINKS_FILE, CUBE_UPDATE_LINKS_FILE):
                    cube.status = 'updating'
                    cube.reboot()
                    shell.log.append('info', f'Cube {cube.ip} - готов к обновлению. Отправлен в перезагрузку.', True)
                else:
                    cube.status = 'skip'
                    shell.log.append('warning', f'Cube {cube.ip} - не готов к обновлению. Пропущен.', True)

#------------------------------------------------------------------------

    for cube in cubes:
        if cube.status == 'updating':
            continue_run = True
    
    if not continue_run:
        addresses.clear()
        cubes.clear()
        os.remove(SHELL_UPDATE_LINKS_FILE)
        shell.log.append('info', 'Нет ни одного готового устройства. Скрипт завершает работу',True)
        return 1

#--- Пауза исполнения скрипта -------------------------------------------
    
    print(f'Делаем паузу на {SCRIPT_PAUSE_BEFORE_VERIFY} минут(ы) перед началом проверки - успешно ли обновление...')
    time.sleep(SCRIPT_PAUSE_BEFORE_VERIFY*60) 

#--- Верификация обновления ---------------------------------------------

    continue_run = False

    for i in range(1, VERIFY_ATTEMPT_MAX_COUNT + 1):
        shell.log.append('info', f'{i}/{VERIFY_ATTEMPT_MAX_COUNT} - Попытка проверки обновления', True)
        for cube in cubes:
            match cube.status:
            
                case 'skip':
                    shell.log.append('warning', f'Cube {cube.ip} - обновление было пропущено', True)
            
                case 'updating':
                    cube.md5_hash = cube.getHash()
                    if cube.md5_hash != cube.md5_hash_prev and cube.uptime() >= CUBE_UPTIME_AFTER:
                        cube.status = 'updated'
                        cube.updater_file = CUBE_UPDATER_FILE_AFTER
                        shell.log.append('info', f'Cube {cube.ip} - обновление прошло успешно', True)
                    else:
                        shell.log.append('error', f'Cube {cube.ip} - факт обновления не зарегистрирован', True)
                        continue_run = True

        if not continue_run:
            break
        print(f'Делаем паузу на {PAUSE_BETWEEN_ATTEMPT} минут(ы) перед следующей попыткой проверки...')
        time.sleep(PAUSE_BETWEEN_ATTEMPT*60)

#--- Очистка от записей и файлов на cube --------------------------------

    for cube in cubes:
        match cube.status:
            case 'updated':
                if cube.isReplaceLine(f'Links_File = "{CUBE_UPDATE_LINKS_FILE}"', 'Links_File = ""'):
                    cube.runCommand(f'rm {CUBE_UPDATE_LINKS_FILE}')

    addresses.clear()
    cubes.clear()
    os.remove(SHELL_UPDATE_LINKS_FILE)
    shell.log.append('info', 'Скрипт завершает работу',True)
    return 0

#------------------------------------------------------------------------

if __name__ == "__main__":
    main()