handOver2/updater.py

# -*- coding: utf-8 -*-
"""
Updater Application
Downloads and extracts updates for the Handover System.
"""

import sys
import os
import time
import zipfile
import argparse
import shutil
import subprocess
import requests
from pathlib import Path
from PySide6.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QLabel, QProgressBar, QMessageBox
)
from PySide6.QtCore import Qt, QThread, Signal

class UpdateWorker(QThread):
    progress_changed = Signal(int)
    status_changed = Signal(str)
    finished = Signal(bool, str)

    def __init__(self, download_url, target_dir, restart_exe):
        super().__init__()
        self.download_url = download_url
        self.target_dir = Path(target_dir)
        self.restart_exe = restart_exe

    def run(self):
        try:
            # 1. Wait for main application to close
            self.status_changed.emit("메인 애플리케이션 종료 대기 중...")
            time.sleep(2)  # Give it a moment
            
            # Simple check: try to rename the executable. If it fails, it's still running.
            # A more robust way would be checking process list, but this is simple and effective.
            exe_path = self.target_dir / self.restart_exe
            max_retries = 10
            for i in range(max_retries):
                try:
                    if exe_path.exists():
                        # Try to open for exclusive access
                        with open(exe_path, 'ab'):
                            pass
                    break
                except IOError:
                    if i == max_retries - 1:
                        raise Exception("애플리케이션이 종료되지 않았습니다. 수동으로 종료해주세요.")
                    time.sleep(1)

            # 2. Download Update
            self.status_changed.emit("업데이트 다운로드 중...")
            local_zip = self.target_dir / "update.zip"
            
            response = requests.get(self.download_url, stream=True)
            response.raise_for_status()
            
            total_size = int(response.headers.get('content-length', 0))
            block_size = 8192
            downloaded = 0
            
            with open(local_zip, 'wb') as f:
                for chunk in response.iter_content(chunk_size=block_size):
                    if chunk:
                        f.write(chunk)
                        downloaded += len(chunk)
                        if total_size > 0:
                            percent = int((downloaded / total_size) * 50) # First 50% is download
                            self.progress_changed.emit(percent)
            
            # 3. Extract Update
            self.status_changed.emit("업데이트 설치 중...")
            with zipfile.ZipFile(local_zip, 'r') as zip_ref:
                file_list = zip_ref.namelist()
                total_files = len(file_list)
                
                for idx, file in enumerate(file_list):
                    zip_ref.extract(file, self.target_dir)
                    percent = 50 + int((idx / total_files) * 50) # Last 50% is extraction
                    self.progress_changed.emit(percent)
            
            # Cleanup
            try:
                os.remove(local_zip)
            except:
                pass
                
            self.progress_changed.emit(100)
            self.status_changed.emit("업데이트 완료!")
            self.finished.emit(True, "성공")
            
        except Exception as e:
            self.finished.emit(False, str(e))

class UpdaterWindow(QWidget):
    def __init__(self, download_url, target_dir, restart_exe):
        super().__init__()
        self.download_url = download_url
        self.target_dir = target_dir
        self.restart_exe = restart_exe
        
        self.init_ui()
        self.start_update()

    def init_ui(self):
        self.setWindowTitle("소프트웨어 업데이트")
        self.setFixedSize(400, 150)
        
        layout = QVBoxLayout()
        
        self.status_label = QLabel("준비 중...")
        self.status_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.status_label)
        
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        layout.addWidget(self.progress_bar)
        
        self.setLayout(layout)

    def start_update(self):
        self.worker = UpdateWorker(self.download_url, self.target_dir, self.restart_exe)
        self.worker.progress_changed.connect(self.progress_bar.setValue)
        self.worker.status_changed.connect(self.status_label.setText)
        self.worker.finished.connect(self.on_finished)
        self.worker.start()

    def on_finished(self, success, message):
        if success:
            # Restart application
            exe_path = Path(self.target_dir) / self.restart_exe
            if exe_path.exists():
                subprocess.Popen([str(exe_path)])
            else:
                QMessageBox.warning(self, "오류", f"실행 파일을 찾을 수 없습니다:\n{exe_path}")
            
            self.close()
        else:
            QMessageBox.critical(self, "업데이트 실패", f"오류가 발생했습니다:\n{message}")
            self.close()

def main():
    parser = argparse.ArgumentParser(description='Handover System Updater')
    parser.add_argument('--url', required=True, help='Download URL for the update')
    parser.add_argument('--target', required=True, help='Target installation directory')
    parser.add_argument('--restart', required=True, help='Executable name to restart')
    
    args = parser.parse_args()
    
    app = QApplication(sys.argv)
    window = UpdaterWindow(args.url, args.target, args.restart)
    window.show()
    sys.exit(app.exec())

if __name__ == "__main__":
    main()