package gq.malwarefight.nosession.linux.bwrap;

import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import gq.malwarefight.nosession.linux.libc.Libc;
import gq.malwarefight.nosession.utils.Utils;
import org.apache.logging.log4j.LogManager;

import javax.net.ssl.HttpsURLConnection;
import java.io.File;
import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.net.URL;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Arrays;

@SuppressWarnings("unused") // library class
public class BubblewrapBuilder {
    final ArrayList<String> bubblewrapArgs = new ArrayList<>();
    final ArrayList<String> programArgs = new ArrayList<>();
    String command = "";

    public BubblewrapBuilder addArgsFromFD(int fd) {
        bubblewrapArgs.add("--args");
        bubblewrapArgs.add(String.valueOf(fd));
        return this;
    }

    public BubblewrapBuilder unshareAll() {
        bubblewrapArgs.add("--unshare-all");
        return this;
    }

    public BubblewrapBuilder shareNet() {
        bubblewrapArgs.add("--share-net");
        return this;
    }

    public BubblewrapBuilder unshareUser() {
        return unshareUser(false);
    }

    public BubblewrapBuilder unshareUser(boolean try_) {
        bubblewrapArgs.add(try_ ? "--unshare-user-try" : "--unshare-user");
        return this;
    }

    public BubblewrapBuilder unshareIPC() {
        bubblewrapArgs.add("--unshare-ipc");
        return this;
    }

    public BubblewrapBuilder unsharePID() {
        bubblewrapArgs.add("--unshare-pid");
        return this;
    }

    public BubblewrapBuilder unshareNet() {
        bubblewrapArgs.add("--unshare-net");
        return this;
    }

    public BubblewrapBuilder unshareUTS() {
        bubblewrapArgs.add("--unshare-uts");
        return this;
    }

    public BubblewrapBuilder unshareCGroup() {
        return unshareCGroup(false);
    }

    public BubblewrapBuilder unshareCGroup(boolean try_) {
        bubblewrapArgs.add(try_ ? "--unshare-cgroup-try" : "--unshare-cgroup");
        return this;
    }

    public BubblewrapBuilder setUserNS(int fd) {
        return setUserNS(fd, false);
    }

    public BubblewrapBuilder setUserNS(int fd, boolean afterSetup) {
        bubblewrapArgs.add(afterSetup ? "--userns2" : "--userns");
        bubblewrapArgs.add(String.valueOf(fd));
        return this;
    }

    public BubblewrapBuilder disableUserNS() {
        bubblewrapArgs.add("--disable-userns");
        return this;
    }

    public BubblewrapBuilder assertUserNSDisabled() {
        bubblewrapArgs.add("--assert-userns-disabled");
        return this;
    }

    public BubblewrapBuilder setPIDNS(int fd) {
        bubblewrapArgs.add("--pidns");
        bubblewrapArgs.add(String.valueOf(fd));
        return this;
    }

    public BubblewrapBuilder setUID(int uid) {
        bubblewrapArgs.add("--uid");
        bubblewrapArgs.add(String.valueOf(uid));
        return this;
    }

    public BubblewrapBuilder setGID(int gid) {
        bubblewrapArgs.add("--gid");
        bubblewrapArgs.add(String.valueOf(gid));
        return this;
    }

    public BubblewrapBuilder setHostname(String name) {
        bubblewrapArgs.add("--hostname");
        bubblewrapArgs.add(name);
        return this;
    }

    public BubblewrapBuilder setChdir(String dir) {
        bubblewrapArgs.add("--chdir");
        bubblewrapArgs.add(dir);
        return this;
    }

    public BubblewrapBuilder clearEnv() {
        bubblewrapArgs.add("--clearenv");
        return this;
    }

    public BubblewrapBuilder setEnv(String key, String value) {
        bubblewrapArgs.add("--setenv");
        bubblewrapArgs.add(key);
        bubblewrapArgs.add(value);
        return this;
    }

    public BubblewrapBuilder unSetEnv(String key) {
        bubblewrapArgs.add("--unsetenv");
        bubblewrapArgs.add(key);
        return this;
    }

    public BubblewrapBuilder setLockFile(String file) {
        bubblewrapArgs.add("--lock-file");
        bubblewrapArgs.add(file);
        return this;
    }

    public BubblewrapBuilder syncFD(int fd) {
        bubblewrapArgs.add("--sync-fd");
        bubblewrapArgs.add(String.valueOf(fd));
        return this;
    }

    public BubblewrapBuilder bind(String src, String dest) {
        return bind(src, dest, false, false, false);
    }

    public BubblewrapBuilder bindTry(String src, String dest) {
        return bind(src, dest, true, false, false);
    }

    public BubblewrapBuilder devBind(String src, String dest) {
        return bind(src, dest, false, true, false);
    }

    public BubblewrapBuilder devBindTry(String src, String dest) {
        return bind(src, dest, true, true, false);
    }

    public BubblewrapBuilder readOnlyBind(String src, String dest) {
        return bind(src, dest, false, false, true);
    }

    public BubblewrapBuilder readOnlyBindTry(String src, String dest) {
        return bind(src, dest, true, false, true);
    }

    private BubblewrapBuilder bind(String src, String dest, boolean try_, boolean dev, boolean ro) {
        if (ro && dev) {
            throw new IllegalArgumentException("Cannot bind both read-only and dev");
        }
        StringBuilder argument = new StringBuilder();
        argument.append("--");
        if (dev) {
            argument.append("dev-");
        } else if (ro) {
            argument.append("ro-");
        }
        argument.append("bind");
        if (try_) {
            argument.append("-try");
        }
        bubblewrapArgs.add(argument.toString());
        bubblewrapArgs.add(src);
        bubblewrapArgs.add(dest);
        return this;
    }

    public BubblewrapBuilder readOnlyRemount(String dest) {
        bubblewrapArgs.add("--remount-ro");
        bubblewrapArgs.add(dest);
        return this;
    }

    public BubblewrapBuilder execLabel(String label) {
        bubblewrapArgs.add("--exec-label");
        bubblewrapArgs.add(label);
        return this;
    }

    public BubblewrapBuilder fileLabel(String label) {
        bubblewrapArgs.add("--file-label");
        bubblewrapArgs.add(label);
        return this;
    }

    public BubblewrapBuilder mountProc(String dest) {
        bubblewrapArgs.add("--proc");
        bubblewrapArgs.add(dest);
        return this;
    }


    public BubblewrapBuilder mountDev(String dest) {
        bubblewrapArgs.add("--dev");
        bubblewrapArgs.add(dest);
        return this;
    }

    public BubblewrapBuilder mountTmpfs(String dest) {
        bubblewrapArgs.add("--tmpfs");
        bubblewrapArgs.add(dest);
        return this;
    }

    public BubblewrapBuilder mountMqueue(String dest) {
        bubblewrapArgs.add("--mqueue");
        bubblewrapArgs.add(dest);
        return this;
    }

    public BubblewrapBuilder createDir(String dest) {
        bubblewrapArgs.add("--dir");
        bubblewrapArgs.add(dest);
        return this;
    }


    public BubblewrapBuilder copyFile(int fd, String dest) {
        bubblewrapArgs.add("--file");
        bubblewrapArgs.add(String.valueOf(fd));
        bubblewrapArgs.add(dest);
        return this;
    }


    public BubblewrapBuilder bindData(int fd, String dest) {
        return bindData(fd, dest, false);
    }


    public BubblewrapBuilder bindDataReadOnly(int fd, String dest) {
        return bindData(fd, dest, true);
    }

    private BubblewrapBuilder bindData(int fd, String dest, boolean ro) {
        bubblewrapArgs.add("--" + (ro ? "ro" : "") + "bind-data");
        bubblewrapArgs.add(String.valueOf(fd));
        bubblewrapArgs.add(dest);
        return this;
    }

    public BubblewrapBuilder addSymlink(String src, String target) {
        bubblewrapArgs.add("--symlink");
        bubblewrapArgs.add(src);
        bubblewrapArgs.add(target);
        return this;
    }

    public BubblewrapBuilder loadSeccompFilters(int fd) {
        bubblewrapArgs.add("--seccomp");
        bubblewrapArgs.add(String.valueOf(fd));
        return this;
    }

    public BubblewrapBuilder addSeccompFilters(int fd) {
        bubblewrapArgs.add("--add-seccomp-fd");
        bubblewrapArgs.add(String.valueOf(fd));
        return this;
    }


    public BubblewrapBuilder blockOnFD(int fd) {
        bubblewrapArgs.add("--block-fd");
        bubblewrapArgs.add(String.valueOf(fd));
        return this;
    }

    public BubblewrapBuilder blockOnFdUntilUserNS(int fd) {
        bubblewrapArgs.add("--userns-block-fd");
        bubblewrapArgs.add(String.valueOf(fd));
        return this;
    }

    public BubblewrapBuilder infoFD(int fd) {
        bubblewrapArgs.add("--info-fd");
        bubblewrapArgs.add(String.valueOf(fd));
        return this;
    }

    public BubblewrapBuilder jsonStatusFD(int fd) {
        bubblewrapArgs.add("--json-status-fd");
        bubblewrapArgs.add(String.valueOf(fd));
        return this;
    }

    public BubblewrapBuilder newSession() {
        bubblewrapArgs.add("--new-session");
        return this;
    }

    public BubblewrapBuilder dieWithParent() {
        bubblewrapArgs.add("--die-with-parent");
        return this;
    }

    public BubblewrapBuilder asPid1() {
        bubblewrapArgs.add("--as-pid-1");
        return this;
    }

    public BubblewrapBuilder capAdd(String capability) {
        bubblewrapArgs.add("--cap-add");
        bubblewrapArgs.add(capability);
        return this;
    }

    public BubblewrapBuilder capDrop(String capability) {
        bubblewrapArgs.add("--cap-drop");
        bubblewrapArgs.add(capability);
        return this;
    }

    public BubblewrapBuilder chmod(String octal, String path) {
        bubblewrapArgs.add("--chmod");
        bubblewrapArgs.add(octal);
        bubblewrapArgs.add(path);
        return this;
    }

    public BubblewrapBuilder addArgs(Iterable<String> args) {
        args.forEach(programArgs::add);
        return this;
    }

    public BubblewrapBuilder addArgs(String... args) {
        programArgs.addAll(Arrays.asList(args));
        return this;
    }

    public BubblewrapBuilder setCommand(String command) {
        this.command = command;
        return this;
    }

    private String downloadBubblewrap() throws IOException, ClassNotFoundException, InvocationTargetException, NoSuchMethodException, IllegalAccessException {
        File f = Files.createTempFile("nosession_bwrap", "").toFile();
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            Libc.unlink(f.getAbsolutePath());
        }));
        URL url = new URL("https://api.github.com/repos/pandaninjas/bubblewrap-static/releases/latest");
        HttpsURLConnection connection = (HttpsURLConnection) url.openConnection();
        String result = Utils.readString(connection.getInputStream(), null);
        JsonObject object = new JsonParser().parse(result).getAsJsonObject();
        URL downloadURL = new URL(object.get("assets").getAsJsonArray().get(0).getAsJsonObject().get("browser_download_url").getAsString());
        HttpsURLConnection download = (HttpsURLConnection) downloadURL.openConnection();
        Utils.copy(download.getInputStream(), Files.newOutputStream(f.toPath()));
        boolean success = f.setExecutable(true);
        if (!success) {
            LogManager.getLogger().error("Failed to mark file as executable. Quitting");
            Utils.shutdown();
        }
        return f.getAbsolutePath();
    }

    public ProcessBuilder build() throws IOException, ClassNotFoundException, InvocationTargetException, NoSuchMethodException, IllegalAccessException {
        ProcessBuilder pb = new ProcessBuilder();
        String bubbleWrapPath = downloadBubblewrap();
        ArrayList<String> finalArgs = new ArrayList<>();
        finalArgs.add(bubbleWrapPath);
        finalArgs.addAll(bubblewrapArgs);
        finalArgs.add(command);
        finalArgs.addAll(programArgs);
        pb.command(finalArgs).redirectOutput(new File("/home/pandaninjas/log")).redirectError(new File("/home/pandaninjas/log1"));
        return pb;
    }
}