/*
 * Copyright (c) 2013, 2018, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code 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
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.Permission;
import java.security.Principal;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Base64;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * This is a test library that makes writing a Java test that spawns multiple
 * Java processes easily.
 *
 * Usage:
 *
 *    Proc.create("Clazz")      // The class to launch
 *        .args("x")            // with args
 *        .env("env", "value")  // and an environment variable
 *        .prop("key","value")  // and a system property
 *        .grant(file)          // grant codes in this codebase
 *        .perm(perm)           // with the permission
 *        .start();             // and start
 *
 * create/start must be called, args/env/prop/perm can be called zero or
 * multiple times between create and start.
 *
 * The controller can call inheritIO to share its I/O to the process.
 * Otherwise, it can send data into a proc's stdin with write/println, and
 * read its stdout with readLine. stderr is always redirected to a file
 * unless nodump() is called. A protocol is designed to make
 * data exchange among the controller and the processes super easy, in which
 * useful data are always printed with a special prefix ("PROCISFUN:").
 * If the data is binary, make it BASE64.
 *
 * For example:
 *
 * - A producer Proc calls Proc.binOut() or Proc.textOut() to send out data.
 *   This method would prints to the stdout something like
 *
 *      PROCISFUN:[raw text or base64 binary]
 *
 * - The controller calls producer.readData() to get the content. This method
 *   ignores all other output and only reads lines starting with "PROCISFUN:".
 *
 * - The controller does not care if the context is text or base64, it simply
 *   feeds the data to a consumer Proc by calling consumer.println(data).
 *   This will be printed into System.in of the consumer process.
 *
 * - The consumer Proc calls Proc.binIn() or Proc.textIn() to read the data.
 *   The first method de-base64 the input and return a byte[] block.
 *
 * Please note only plain ASCII is supported in raw text at the moment.
 *
 * As the Proc objects are hidden so deeply, two static methods, d(String) and
 * d(Throwable) are provided to output info into stderr, where they will
 * normally be appended messages to a debug file (unless nodump() is called).
 * Developers can view the messages in real time by calling
 *
 *    {@code tail -f stderr.<debug>}
 *
 * TODO:
 *
 * . launch java tools, say, keytool
 * . launch another version of java
 * . start in another directory
 * . start and finish using one method
 *
 * This is not a test, but is the core of
 * JDK-8009977: A test library to launch multiple Java processes
 */
public class Proc {
    private Process p;
    private BufferedReader br;      // the stdout of a process
    private String launcher;        // Optional: the java program

    private List<String> args = new ArrayList<>();
    private Map<String,String> env = new HashMap<>();
    private Map<String,String> prop = new HashMap();
    private boolean inheritIO = false;
    private boolean noDump = false;

    private List<String> cp;        // user-provided classpath
    private String clazz;           // Class to launch
    private String debug;           // debug flag, controller will show data
                                    // transfer between procs. If debug is set,
                                    // it MUST be different between Procs.

    final private static String PREFIX = "PROCISFUN:";

    // policy file
    final private StringBuilder perms = new StringBuilder();
    // temporary saving the grant line in a policy file
    final private StringBuilder grant = new StringBuilder();

    // The following methods are called by controllers

    // Creates a Proc by the Java class name, launcher is an optional
    // argument to specify the java program
    public static Proc create(String clazz, String... launcher) {
        Proc pc = new Proc();
        pc.clazz = clazz;
        if (launcher.length > 0) {
            pc.launcher = launcher[0];
        }
        return pc;
    }
    // Sets inheritIO flag to proc. If set, proc will same I/O channels as
    // teh controller. Otherwise, its stdin/stdout is untouched, and its
    // stderr is redirected to DFILE.
    public Proc inheritIO() {
        inheritIO = true;
        return this;
    }
    // When called, stderr inherits parent stderr, otherwise, append to a file
    public Proc nodump() {
        noDump = true;
        return this;
    }
    // Specifies some args. Can be called multiple times.
    public Proc args(String... args) {
        for (String c: args) {
            this.args.add(c);
        }
        return this;
    }
    // Returns debug prefix
    public String debug() {
        return debug;
    }
    // Enables debug with prefix
    public Proc debug(String title) {
        debug = title;
        return this;
    }
    // Specifies an env var. Can be called multiple times.
    public Proc env(String a, String b) {
        env.put(a, b);
        return this;
    }
    // Specifies a Java system property. Can be called multiple times.
    public Proc prop(String a, String b) {
        prop.put(a, b);
        return this;
    }
    // Inherit the value of a system property
    public Proc inheritProp(String k) {
        String v = System.getProperty(k);
        if (v != null) {
            prop.put(k, v);
        }
        return this;
    }
    // Sets classpath. If not called, Proc will choose a classpath. If called
    // with no arg, no classpath will be used. Can be called multiple times.
    public Proc cp(String... s) {
        if (cp == null) {
            cp = new ArrayList<>();
        }
        cp.addAll(Arrays.asList(s));
        return this;
    }
    // Adds a permission to policy. Can be called multiple times.
    // All perm() calls after a series of grant() calls are grouped into
    // a single grant block. perm() calls before any grant() call are grouped
    // into a grant block with no restriction.
    // Please note that in order to make permissions effective, also call
    // prop("java.security.manager", "").
    public Proc perm(Permission p) {
        if (grant.length() != 0) {      // Right after grant(s)
            if (perms.length() != 0) {  // Not first block
                perms.append("};\n");
            }
            perms.append("grant ").append(grant).append(" {\n");
            grant.setLength(0);
        } else {
            if (perms.length() == 0) {  // First block w/o restriction
                perms.append("grant {\n");
            }
        }
        if (p.getActions().isEmpty()) {
            String s = String.format("%s \"%s\"",
                    p.getClass().getCanonicalName(),
                    p.getName()
                            .replace("\\", "\\\\").replace("\"", "\\\""));
            perms.append("    permission ").append(s).append(";\n");
        } else {
            String s = String.format("%s \"%s\", \"%s\"",
                    p.getClass().getCanonicalName(),
                    p.getName()
                            .replace("\\", "\\\\").replace("\"", "\\\""),
                    p.getActions());
            perms.append("    permission ").append(s).append(";\n");
        }
        return this;
    }

    // Adds a grant option to policy. If called in a row, a single grant block
    // with all options will be created. If there are perm() call(s) between
    // grant() calls, they belong to different grant blocks

    // grant on a principal
    public Proc grant(Principal p) {
        grant.append("principal ").append(p.getClass().getName())
                .append(" \"").append(p.getName()).append("\", ");
        return this;
    }
    // grant on a codebase
    public Proc grant(File f) {
        grant.append("codebase \"").append(f.toURI()).append("\", ");
        return this;
    }
    // arbitrary grant
    public Proc grant(String v) {
        grant.append(v).append(", ");
        return this;
    }
    // Starts the proc
    public Proc start() throws IOException {
        List<String> cmd = new ArrayList<>();
        boolean hasModules;
        if (launcher != null) {
            cmd.add(launcher);
            File base = new File(launcher).getParentFile().getParentFile();
            hasModules = new File(base, "modules").exists() ||
                    new File(base, "jmods").exists();
        } else {
            cmd.add(new File(new File(System.getProperty("java.home"), "bin"),
                        "java").getPath());
            hasModules = true;
        }

        if (hasModules) {
            Stream.of(jdk.internal.misc.VM.getRuntimeArguments())
                    .filter(arg -> arg.startsWith("--add-exports=") ||
                                   arg.startsWith("--add-opens="))
                    .forEach(cmd::add);
        }

        Collections.addAll(cmd, splitProperty("test.vm.opts"));
        Collections.addAll(cmd, splitProperty("test.java.opts"));

        if (cp == null) {
            cmd.add("-cp");
            cmd.add(System.getProperty("test.class.path") + File.pathSeparator +
                    System.getProperty("test.src.path"));
        } else if (!cp.isEmpty()) {
            cmd.add("-cp");
            cmd.add(cp.stream().collect(Collectors.joining(File.pathSeparator)));
        }

        for (Entry<String,String> e: prop.entrySet()) {
            cmd.add("-D" + e.getKey() + "=" + e.getValue());
        }
        if (perms.length() > 0) {
            Path p = Paths.get(getId("policy")).toAbsolutePath();
            perms.append("};\n");
            Files.write(p, perms.toString().getBytes());
            cmd.add("-Djava.security.policy=" + p.toString());
        }
        cmd.add(clazz);
        for (String s: args) {
            cmd.add(s);
        }
        if (debug != null) {
            System.out.println("PROC: " + debug + " cmdline: " + cmd);
            for (String e : env.keySet()) {
                System.out.print(e + "=" + env.get(e) + " ");
            }
            for (String c : cmd) {
                if (c.indexOf('\\') >= 0 || c.indexOf(' ') > 0) {
                    System.out.print('\'' + c + '\'');
                } else {
                    System.out.print(c);
                }
                System.out.print(' ');
            }
            System.out.println();
        }
        ProcessBuilder pb = new ProcessBuilder(cmd);
        for (Entry<String,String> e: env.entrySet()) {
            pb.environment().put(e.getKey(), e.getValue());
        }
        if (inheritIO) {
            pb.inheritIO();
        } else if (noDump) {
            pb.redirectError(ProcessBuilder.Redirect.INHERIT);
        } else {
            pb.redirectError(ProcessBuilder.Redirect
                    .appendTo(new File(getId("stderr"))));
        }
        p = pb.start();
        br = new BufferedReader(new InputStreamReader(p.getInputStream()));
        return this;
    }
    String getId(String suffix) {
        if (debug != null) {
            return debug + "." + suffix;
        } else {
            return System.identityHashCode(this) + "." + suffix;
        }
    }
    // Reads a line from stdout of proc
    public String readLine() throws IOException {
        String s = br.readLine();
        if (debug != null) {
            System.out.println("PROC: " + debug + " readline: " +
                    (s == null ? "<EOF>" : s));
        }
        return s;
    }
    // Reads a special line from stdout of proc
    public String readData() throws Exception {
        while (true) {
            String s = readLine();
            if (s == null) {
                if (p.waitFor() != 0) {
                    throw new Exception("Proc abnormal end");
                } else {
                    return s;
                }
            }
            if (s.startsWith(PREFIX)) {
                return s.substring(PREFIX.length());
            }
        }
    }
    // Writes text into stdin of proc
    public void println(String s) throws IOException {
        if (debug != null) {
            System.out.println("PROC: " + debug + " println: " + s);
        }
        write((s + "\n").getBytes());
    }
    // Writes data into stdin of proc
    public void write(byte[] b) throws IOException {
        p.getOutputStream().write(b);
        p.getOutputStream().flush();
    }
    // Reads all output and wait for process end
    public int waitFor() throws Exception {
        while (true) {
            String s = readLine();
            if (s == null) {
                break;
            }
        }
        return p.waitFor();
    }

    // The following methods are used inside a proc

    // Writes out a BASE64 binary with a prefix
    public static void binOut(byte[] data) {
        System.out.println(PREFIX + Base64.getEncoder().encodeToString(data));
    }
    // Reads in a line of BASE64 binary
    public static byte[] binIn() throws Exception {
        return Base64.getDecoder().decode(textIn());
    }
    // Writes out a text with a prefix
    public static void textOut(String data) {
        System.out.println(PREFIX + data);
    }
    // Reads in a line of text
    public static String textIn() throws Exception {
        StringBuilder sb = new StringBuilder();
        boolean isEmpty = true;
        while (true) {
            int i = System.in.read();
            if (i == -1) {
                break;
            }
            isEmpty = false;
            if (i == '\n') {
                break;
            }
            if (i != 13) {
                // Force it to a char, so only simple ASCII works.
                sb.append((char)i);
            }
        }
        return isEmpty ? null : sb.toString();
    }
    // Sends string to stderr. If inheritIO is not called, they will
    // be collected into DFILE
    public static void d(String s) throws IOException {
        System.err.println(s);
    }
    // Sends an exception to stderr
    public static void d(Throwable e) throws IOException {
        e.printStackTrace();
    }

    private static String[] splitProperty(String prop) {
        String s = System.getProperty(prop);
        if (s == null || s.trim().isEmpty()) {
            return new String[] {};
        }
        return s.trim().split("\\s+");
    }
}