/*
 * This file is part of LSPosed.
 *
 * LSPosed is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * LSPosed is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with LSPosed.  If not, see <https://www.gnu.org/licenses/>.
 *
 * Copyright (C) 2022 LSPosed Contributors
 */

package org.lsposed.lspd.service;

import static org.lsposed.lspd.ILSPManagerService.DEX2OAT_CRASHED;
import static org.lsposed.lspd.ILSPManagerService.DEX2OAT_MOUNT_FAILED;
import static org.lsposed.lspd.ILSPManagerService.DEX2OAT_OK;
import static org.lsposed.lspd.ILSPManagerService.DEX2OAT_SELINUX_PERMISSIVE;
import static org.lsposed.lspd.ILSPManagerService.DEX2OAT_SEPOLICY_INCORRECT;

import android.net.LocalServerSocket;
import android.os.Build;
import android.os.FileObserver;
import android.os.Process;
import android.os.SELinux;
import android.system.ErrnoException;
import android.system.Os;
import android.system.OsConstants;
import android.util.Log;

import androidx.annotation.Nullable;
import androidx.annotation.RequiresApi;

import java.io.File;
import java.io.FileDescriptor;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;

@RequiresApi(Build.VERSION_CODES.Q)
public class Dex2OatService implements Runnable {
    private static final String TAG = "LSPosedDex2Oat";
    private static final String WRAPPER32 = "bin/dex2oat32";
    private static final String WRAPPER64 = "bin/dex2oat64";
    private static final String HOOKER32 = "bin/liboat_hook32.so";
    private static final String HOOKER64 = "bin/liboat_hook64.so";

    private final String[] dex2oatArray = new String[6];
    private final FileDescriptor[] fdArray = new FileDescriptor[6];
    private final FileObserver selinuxObserver;
    private int compatibility = DEX2OAT_OK;

    private void openDex2oat(int id, String path) {
        try {
            var fd = Os.open(path, OsConstants.O_RDONLY, 0);
            dex2oatArray[id] = path;
            fdArray[id] = fd;
        } catch (ErrnoException ignored) {
        }
    }

    /**
     * Checks the ELF header of the target file.
     * If 32-bit -> Assigns to Index 0 (Release) or 1 (Debug).
     * If 64-bit -> Assigns to Index 2 (Release) or 3 (Debug).
     */
    private void checkAndAddDex2Oat(String path) {
        if (path == null)
            return;
        File file = new File(path);
        if (!file.exists())
            return;

        try (FileInputStream fis = new FileInputStream(file)) {
            byte[] header = new byte[5];
            if (fis.read(header) != 5)
                return;

            // 1. Verify ELF Magic: 0x7F 'E' 'L' 'F'
            if (header[0] != 0x7F || header[1] != 'E' || header[2] != 'L' || header[3] != 'F') {
                return;
            }

            // 2. Check Architecture (header[4]): 1 = 32-bit, 2 = 64-bit
            boolean is32Bit = (header[4] == 1);
            boolean is64Bit = (header[4] == 2);
            boolean isDebug = path.contains("dex2oatd");

            int index = -1;

            if (is32Bit) {
                index = isDebug ? 1 : 0; // Index 0/1 maps to r32/d32 in C++
            } else if (is64Bit) {
                index = isDebug ? 3 : 2; // Index 2/3 maps to r64/d64 in C++
            }

            // 3. Assign to the detected slot
            if (index != -1 && dex2oatArray[index] == null) {
                dex2oatArray[index] = path;
                try {
                    // Open the FD for the wrapper to use later
                    fdArray[index] = Os.open(path, OsConstants.O_RDONLY, 0);
                    Log.i(TAG, "Detected " + path + " as " + (is64Bit ? "64-bit" : "32-bit") + " -> Assigned Index "
                            + index);
                } catch (ErrnoException e) {
                    Log.e(TAG, "Failed to open FD for " + path, e);
                    dex2oatArray[index] = null;
                }
            }
        } catch (IOException e) {
            // File not readable, skip
        }
    }

    public Dex2OatService() {
        if (Build.VERSION.SDK_INT == Build.VERSION_CODES.Q) {
            // Android 10: Check the standard path.
            // Logic will detect if it is 32-bit and put it in Index 0.
            checkAndAddDex2Oat("/apex/com.android.runtime/bin/dex2oat");
            checkAndAddDex2Oat("/apex/com.android.runtime/bin/dex2oatd");

            // Check for explicit 64-bit paths (just in case)
            checkAndAddDex2Oat("/apex/com.android.runtime/bin/dex2oat64");
            checkAndAddDex2Oat("/apex/com.android.runtime/bin/dex2oatd64");
        } else {
            checkAndAddDex2Oat("/apex/com.android.art/bin/dex2oat32");
            checkAndAddDex2Oat("/apex/com.android.art/bin/dex2oatd32");
            checkAndAddDex2Oat("/apex/com.android.art/bin/dex2oat64");
            checkAndAddDex2Oat("/apex/com.android.art/bin/dex2oatd64");
        }

        openDex2oat(4, "/data/adb/modules/zygisk_lsposed/bin/liboat_hook32.so");
        openDex2oat(5, "/data/adb/modules/zygisk_lsposed/bin/liboat_hook64.so");

        var enforce = Paths.get("/sys/fs/selinux/enforce");
        var policy = Paths.get("/sys/fs/selinux/policy");
        var list = new ArrayList<File>();
        list.add(enforce.toFile());
        list.add(policy.toFile());
        selinuxObserver = new FileObserver(list, FileObserver.CLOSE_WRITE) {
            @Override
            public synchronized void onEvent(int i, @Nullable String s) {
                Log.d(TAG, "SELinux status changed");
                if (compatibility == DEX2OAT_CRASHED) {
                    stopWatching();
                    return;
                }

                boolean enforcing = false;
                try (var is = Files.newInputStream(enforce)) {
                    enforcing = is.read() == '1';
                } catch (IOException ignored) {
                }

                if (!enforcing) {
                    if (compatibility == DEX2OAT_OK) doMount(false);
                    compatibility = DEX2OAT_SELINUX_PERMISSIVE;
                } else if (SELinux.checkSELinuxAccess("u:r:untrusted_app:s0",
                        "u:object_r:dex2oat_exec:s0", "file", "execute")
                        || SELinux.checkSELinuxAccess("u:r:untrusted_app:s0",
                        "u:object_r:dex2oat_exec:s0", "file", "execute_no_trans")) {
                    if (compatibility == DEX2OAT_OK) doMount(false);
                    compatibility = DEX2OAT_SEPOLICY_INCORRECT;
                } else if (compatibility != DEX2OAT_OK) {
                    doMount(true);
                    if (notMounted()) {
                        doMount(false);
                        compatibility = DEX2OAT_MOUNT_FAILED;
                        stopWatching();
                    } else {
                        compatibility = DEX2OAT_OK;
                    }
                }
            }

            @Override
            public void stopWatching() {
                super.stopWatching();
                Log.w(TAG, "SELinux observer stopped");
            }
        };
    }

    private boolean notMounted() {
        for (int i = 0; i < dex2oatArray.length && i < 4; i++) {
            var bin = dex2oatArray[i];
            if (bin == null) continue;
            try {
                var apex = Os.stat("/proc/1/root" + bin);
                var wrapper = Os.stat(i < 2 ? WRAPPER32 : WRAPPER64);
                if (apex.st_dev != wrapper.st_dev || apex.st_ino != wrapper.st_ino) {
                    Log.w(TAG, "Check mount failed for " + bin);
                    return true;
                }
            } catch (ErrnoException e) {
                Log.e(TAG, "Check mount failed for " + bin, e);
                return true;
            }
        }
        Log.d(TAG, "Check mount succeeded");
        return false;
    }

    private void doMount(boolean enabled) {
        doMountNative(enabled, dex2oatArray[0], dex2oatArray[1], dex2oatArray[2], dex2oatArray[3]);
    }

    public void start() {
        if (notMounted()) { // Already mounted when restart daemon
            doMount(true);
            if (notMounted()) {
                doMount(false);
                compatibility = DEX2OAT_MOUNT_FAILED;
                return;
            }
        }

        var thread = new Thread(this);
        thread.setName("dex2oat");
        thread.start();
        selinuxObserver.startWatching();
        selinuxObserver.onEvent(0, null);
    }

    @Override
    public void run() {
        Log.i(TAG, "Dex2oat wrapper daemon start");
        var sockPath = getSockPath();
        Log.d(TAG, "wrapper path: " + sockPath);
        var xposed_file = "u:object_r:xposed_file:s0";
        var dex2oat_exec = "u:object_r:dex2oat_exec:s0";
        if (SELinux.checkSELinuxAccess("u:r:dex2oat:s0", dex2oat_exec,
                "file", "execute_no_trans")) {
            SELinux.setFileContext(WRAPPER32, dex2oat_exec);
            SELinux.setFileContext(WRAPPER64, dex2oat_exec);
            setSockCreateContext("u:r:dex2oat:s0");
        } else {
            SELinux.setFileContext(WRAPPER32, xposed_file);
            SELinux.setFileContext(WRAPPER64, xposed_file);
            setSockCreateContext("u:r:installd:s0");
        }
        SELinux.setFileContext(HOOKER32, xposed_file);
        SELinux.setFileContext(HOOKER64, xposed_file);
        try (var server = new LocalServerSocket(sockPath)) {
            setSockCreateContext(null);
            while (true) {
                try (var client = server.accept();
                     var is = client.getInputStream();
                     var os = client.getOutputStream()) {
                    var id = is.read();
                    var fd = new FileDescriptor[]{fdArray[id]};
                    client.setFileDescriptorsForSend(fd);
                    os.write(1);
                    Log.d(TAG, "Sent fd of " + dex2oatArray[id]);
                }
            }
        } catch (IOException e) {
            Log.e(TAG, "Dex2oat wrapper daemon crashed", e);
            if (compatibility == DEX2OAT_OK) {
                doMount(false);
                compatibility = DEX2OAT_CRASHED;
            }
        }
    }

    public int getCompatibility() {
        return compatibility;
    }

    private native void doMountNative(boolean enabled,
                                      String r32, String d32, String r64, String d64);

    private static native boolean setSockCreateContext(String context);

    private native String getSockPath();
}