mesclient-OKNG/src/com/mes/update/UpdateInstaller.java

package com.mes.update;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.lang.management.ManagementFactory;
import java.nio.charset.Charset;
import java.util.Enumeration;
import java.util.zip.ZipEntry;
import java.util.zip.ZipFile;
import java.util.zip.ZipInputStream;

/**
 * 解压升级包并生成 Windows 替换脚本（进程退出后覆盖文件并重启）。
 */
public class UpdateInstaller {

    private static final Logger log = LoggerFactory.getLogger(UpdateInstaller.class);

    /**
     * 解压 zip 到目标目录（已存在则清空重建）。
     */
    public File unzip(File zipFile, File extractDir) throws Exception {
        if (extractDir.exists()) {
            deleteRecursively(extractDir);
        }
        if (!extractDir.mkdirs()) {
            throw new IllegalStateException("无法创建解压目录: " + extractDir.getAbsolutePath());
        }
        try {
            unzipWithZipFile(zipFile, extractDir);
        } catch (Exception encodingIssue) {
            log.warn("ZipFile 解压失败，回退 ZipInputStream：{}", encodingIssue.getMessage());
            if (extractDir.exists()) {
                deleteRecursively(extractDir);
            }
            if (!extractDir.mkdirs()) {
                throw new IllegalStateException("无法创建解压目录: " + extractDir.getAbsolutePath());
            }
            unzipFallback(zipFile, extractDir);
        }
        return flattenIfSingleRoot(extractDir);
    }

    private void unzipWithZipFile(File zipFile, File extractDir) throws Exception {
        try (ZipFile zip = new ZipFile(zipFile, Charset.forName("UTF-8"))) {
            Enumeration<? extends ZipEntry> entries = zip.entries();
            byte[] buffer = new byte[8192];
            while (entries.hasMoreElements()) {
                ZipEntry entry = entries.nextElement();
                writeZipEntry(extractDir, entry, zip.getInputStream(entry), buffer);
            }
        }
    }

    private void unzipFallback(File zipFile, File extractDir) throws Exception {
        byte[] buffer = new byte[8192];
        try (ZipInputStream zis = new ZipInputStream(new FileInputStream(zipFile))) {
            ZipEntry entry;
            while ((entry = zis.getNextEntry()) != null) {
                writeZipEntry(extractDir, entry, zis, buffer);
                zis.closeEntry();
            }
        }
    }

    private void writeZipEntry(File extractDir, ZipEntry entry, java.io.InputStream in, byte[] buffer) throws Exception {
        File outFile = new File(extractDir, entry.getName());
        if (!isSafeExtractPath(extractDir, outFile)) {
            throw new IllegalStateException("非法压缩条目: " + entry.getName());
        }
        if (entry.isDirectory()) {
            if (!outFile.exists() && !outFile.mkdirs()) {
                throw new IllegalStateException("无法创建目录: " + outFile.getAbsolutePath());
            }
            return;
        }
        File parent = outFile.getParentFile();
        if (parent != null && !parent.exists() && !parent.mkdirs()) {
            throw new IllegalStateException("无法创建目录: " + parent.getAbsolutePath());
        }
        try (FileOutputStream out = new FileOutputStream(outFile);
             BufferedOutputStream bos = new BufferedOutputStream(out)) {
            int n;
            while ((n = in.read(buffer)) >= 0) {
                bos.write(buffer, 0, n);
            }
        }
    }

    private File flattenIfSingleRoot(File extractDir) throws Exception {
        File[] children = extractDir.listFiles();
        if (children == null || children.length != 1 || !children[0].isDirectory()) {
            return extractDir;
        }
        File nested = children[0];
        if (!hasAppPayload(nested)) {
            return extractDir;
        }
        File flat = new File(extractDir.getParentFile(), extractDir.getName() + "-flat");
        if (flat.exists()) {
            deleteRecursively(flat);
        }
        if (!nested.renameTo(flat)) {
            copyDirectory(nested, flat);
            deleteRecursively(extractDir);
            return flat;
        }
        deleteRecursively(extractDir);
        return flat;
    }

    private boolean hasAppPayload(File dir) {
        File[] files = dir.listFiles();
        if (files == null) {
            return false;
        }
        for (File f : files) {
            String name = f.getName().toLowerCase();
            if (name.endsWith(".jar") || name.endsWith(".exe") || name.equals("app.version")) {
                return true;
            }
        }
        return false;
    }

    /**
     * 生成并启动 update.bat：等待当前 PID 退出后覆盖安装目录并启动 EXE。
     */
    public void launchReplaceScript(File installDir, File payloadDir, String exeName) throws Exception {
        if (installDir == null || !installDir.isDirectory()) {
            throw new IllegalStateException("安装目录无效");
        }
        if (payloadDir == null || !payloadDir.isDirectory()) {
            throw new IllegalStateException("升级包内容目录无效");
        }
        String exe = (exeName == null || exeName.trim().isEmpty()) ? resolveExeName(installDir) : exeName.trim();
        long pid = currentPid();
        File script = new File(System.getProperty("java.io.tmpdir"), "mesclient-update-" + pid + ".bat");
        writeBat(script, pid, installDir, payloadDir, exe);

        ProcessBuilder pb = new ProcessBuilder(
                "cmd.exe", "/c", "start", "\"mesclient-update\"", "/min", script.getAbsolutePath());
        pb.directory(installDir);
        pb.start();
        log.info("已启动升级脚本：{}", script.getAbsolutePath());
    }

    private void writeBat(File script, long pid, File installDir, File payloadDir, String exeName) throws Exception {
        Charset charset = Charset.defaultCharset();
        try (PrintWriter out = new PrintWriter(new OutputStreamWriter(new FileOutputStream(script), charset))) {
            out.println("@echo off");
            out.println("setlocal EnableExtensions");
            out.println("set \"PID=" + pid + "\"");
            out.println("set \"INSTALL_DIR=" + installDir.getAbsolutePath() + "\"");
            out.println("set \"SOURCE_DIR=" + payloadDir.getAbsolutePath() + "\"");
            out.println("set \"EXE_NAME=" + exeName + "\"");
            out.println("echo [mes-update] waiting process %PID% ...");
            out.println(":wait");
            out.println("tasklist /FI \"PID eq %PID%\" 2>NUL | findstr /I /C:\" %PID% \" >NUL");
            out.println("if not errorlevel 1 (");
            out.println("  timeout /t 1 /nobreak >NUL");
            out.println("  goto wait");
            out.println(")");
            out.println("echo [mes-update] copying files ...");
            out.println("xcopy /E /Y /I /Q \"%SOURCE_DIR%\\*\" \"%INSTALL_DIR%\\\" >NUL");
            out.println("if errorlevel 1 (");
            out.println("  echo [mes-update] copy failed");
            out.println("  pause");
            out.println("  exit /b 1");
            out.println(")");
            out.println("echo [mes-update] starting %EXE_NAME%");
            out.println("start \"\" \"%INSTALL_DIR%\\%EXE_NAME%\"");
            out.println("rmdir /S /Q \"%SOURCE_DIR%\" >NUL 2>&1");
            out.println("del \"%~f0\" >NUL 2>&1");
            out.println("endlocal");
            out.println("exit /b 0");
        }
    }

    public static long currentPid() {
        try {
            String name = ManagementFactory.getRuntimeMXBean().getName();
            int at = name.indexOf('@');
            if (at > 0) {
                return Long.parseLong(name.substring(0, at));
            }
        } catch (Exception ignored) {
        }
        return 0L;
    }

    public static String resolveExeName(File installDir) {
        File preferred = new File(installDir, "MesClient.exe");
        if (preferred.isFile()) {
            return preferred.getName();
        }
        File[] files = installDir.listFiles();
        if (files != null) {
            for (File f : files) {
                if (f.isFile() && f.getName().toLowerCase().endsWith(".exe")) {
                    return f.getName();
                }
            }
        }
        return "MesClient.exe";
    }

    public static File resolveInstallDir() {
        try {
            java.net.URL loc = UpdateInstaller.class.getProtectionDomain().getCodeSource().getLocation();
            File code = new File(loc.toURI());
            if (code.isFile()) {
                return code.getParentFile();
            }
            return new File(System.getProperty("user.dir"));
        } catch (Exception e) {
            return new File(System.getProperty("user.dir"));
        }
    }

    private static boolean isSafeExtractPath(File baseDir, File target) throws Exception {
        String base = baseDir.getCanonicalPath();
        String path = target.getCanonicalPath();
        return path.startsWith(base + File.separator) || path.equals(base);
    }

    private static void copyDirectory(File src, File dest) throws Exception {
        if (!dest.exists() && !dest.mkdirs()) {
            throw new IllegalStateException("无法创建目录: " + dest.getAbsolutePath());
        }
        File[] files = src.listFiles();
        if (files == null) {
            return;
        }
        byte[] buffer = new byte[8192];
        for (File file : files) {
            File target = new File(dest, file.getName());
            if (file.isDirectory()) {
                copyDirectory(file, target);
            } else {
                try (FileInputStream in = new FileInputStream(file);
                     FileOutputStream out = new FileOutputStream(target)) {
                    int n;
                    while ((n = in.read(buffer)) >= 0) {
                        out.write(buffer, 0, n);
                    }
                }
            }
        }
    }

    public static void deleteRecursively(File file) {
        if (file == null || !file.exists()) {
            return;
        }
        if (file.isDirectory()) {
            File[] children = file.listFiles();
            if (children != null) {
                for (File child : children) {
                    deleteRecursively(child);
                }
            }
        }
        // noinspection ResultOfMethodCallIgnored
        file.delete();
    }
}