/*

 * To change this license header, choose License Headers in Project Properties.

 * To change this template file, choose Tools | Templates

 * and open the template in the editor.

 */

package javax.management.remote.rest;

import com.oracle.jmx.remote.rest.http.PostRequestHandler;

import com.oracle.jmx.remote.rest.http.GetRequestHandler;

import com.sun.jmx.remote.security.JMXPluggableAuthenticator;

import com.sun.jmx.remote.security.JMXSubjectDomainCombiner;

import com.sun.jmx.remote.security.SubjectDelegator;

import com.sun.net.httpserver.*;

import javax.management.*;

import javax.management.relation.MBeanServerNotificationFilter;

import javax.management.remote.JMXAuthenticator;

import com.oracle.jmx.remote.rest.json.JSONArray;

import com.oracle.jmx.remote.rest.json.JSONObject;

import com.oracle.jmx.remote.rest.json.JSONPrimitive;

import com.oracle.jmx.remote.rest.mapper.JSONMappingFactory;

import javax.security.auth.Subject;

import java.io.*;

import java.lang.reflect.InvocationHandler;

import java.lang.reflect.Method;

import java.lang.reflect.Proxy;

import java.net.*;

import java.security.AccessControlContext;

import java.security.AccessController;

import java.security.PrivilegedAction;

import java.util.*;

import com.sun.net.httpserver.Authenticator;

/**

 * @author harsha

 */

public final class JmxRestAdapterImpl implements JmxRestAdapter, NotificationListener {

    // TODO: These should be wrapped in ReadWriteLock

    private static final Map<String, MBeanServer> authMBeanServer = new HashMap<>();

    final List<String> allowedMbeans = new ArrayList<>();

    private final HttpServer httpServer;

    private final String contextStr;

    private final Map<String, ?> env;

    private final MBeanServer mbeanServer;

    private final String realm = "";

    private final GetRequestHandler getHandler;

    private final PostRequestHandler postHandler;

    private HttpContext httpContext;

    private JMXAuthenticator authenticator = null;

    private boolean started = false;

    JmxRestAdapterImpl(HttpServer hServer, String context, Map<String, ?> env, MBeanServer mbeanServer) {

        httpServer = hServer;

        this.contextStr = context;

        this.env = env;

        this.mbeanServer = mbeanServer;

        if (env.get("jmx.remote.x.authentication") != null) {

            authenticator = (JMXAuthenticator) env.get(JmxRestAdapterImpl.AUTHENTICATOR);

            if (authenticator == null) {

                if (env.get("jmx.remote.x.password.file") != null

                        || env.get("jmx.remote.x.login.config") != null) {

                    authenticator = new JMXPluggableAuthenticator(env);

                } else {

                    // Throw exception for invalid authentication config

                }

            }

        }

        introspectMBeanTypes(mbeanServer);

        MBeanServerNotificationFilter filter = new MBeanServerNotificationFilter();

        filter.enableAllObjectNames();

        try {

            mbeanServer.addNotificationListener(MBeanServerDelegate.DELEGATE_NAME, this, filter, null);

        } catch (InstanceNotFoundException ex) {

        }

        getHandler = new GetRequestHandler(mbeanServer, allowedMbeans);

        postHandler = new PostRequestHandler(mbeanServer, allowedMbeans);

    }

    @Override

    public synchronized void start() {

        if (!started) {

            httpContext = httpServer.createContext("/jmx/" + contextStr + "/", new DefaultHTTPHandler());

            if (env.get("jmx.remote.x.authentication") != null) {

                httpContext.setAuthenticator(new BasicAuthenticator());

            }

            started = true;

        }

    }

    @Override

    public synchronized void stop() {

        if (!started) {

            throw new IllegalStateException("Rest Adapter not started yet");

        }

        httpServer.removeContext(httpContext);

        started = false;

    }

    @Override

    public String getBaseUrl() {

        if (!started) {

            throw new IllegalStateException("Adapter not started");

        }

        try {

            if (httpServer instanceof HttpsServer) {

                return "https://" + InetAddress.getLocalHost().getHostName() + ":" + httpServer.getAddress().getPort() + "/jmx/" + contextStr + "/";

            }

            return "http://" + InetAddress.getLocalHost().getHostName() + ":" + httpServer.getAddress().getPort() + "/jmx/" + contextStr + "/";

        } catch (UnknownHostException ex) {

            return "http://localhost" + ":" + httpServer.getAddress().getPort() + "/jmx/" + contextStr + "/";

        }

    }

    @Override

    public MBeanServer getMBeanServer() {

        return mbeanServer;

    }

    private boolean isMBeanAllowed(ObjectName objName) {

        try {

            MBeanInfo mInfo = mbeanServer.getMBeanInfo(objName);

            MBeanAttributeInfo[] attrsInfo = mInfo.getAttributes();

            for (MBeanAttributeInfo attrInfo : attrsInfo) {

                String type = attrInfo.getType();

                if (!JSONMappingFactory.INSTANCE.isTypeMapped(type)) {

                    return false;

                }

            }

            MBeanOperationInfo[] operations = mInfo.getOperations();

            for (MBeanOperationInfo opInfo : operations) {

                MBeanParameterInfo[] signature = opInfo.getSignature();

                for (MBeanParameterInfo sig : signature) {

                    if (!JSONMappingFactory.INSTANCE.isTypeMapped(sig.getType())) {

                        return false;

                    }

                }

                if (!JSONMappingFactory.INSTANCE.isTypeMapped(opInfo.getReturnType())) {

                    return false;

                }

            }

            return true;

        } catch (InstanceNotFoundException | IntrospectionException | ReflectionException | ClassNotFoundException ex) {

            ex.printStackTrace();

            return false;

        }

    }

    private void introspectMBeanTypes(MBeanServer server) {

        if (allowedMbeans.isEmpty()) {

            Set<ObjectInstance> allMBeans = server.queryMBeans(null, null); // get all Mbeans

            allMBeans.stream().filter((objIns) -> (isMBeanAllowed(objIns.getObjectName())))

                    .forEachOrdered(objIns -> allowedMbeans.add(objIns.getObjectName().toString()));

        }

    }

    @Override

    public void handleNotification(Notification notification, Object handback) {

        try {

            MBeanServerNotification mbs = (MBeanServerNotification) notification;

            if (MBeanServerNotification.REGISTRATION_NOTIFICATION.equals(mbs.getType())) {

                ObjectName mBeanName = mbs.getMBeanName();

                if (isMBeanAllowed(mBeanName)) {

                    allowedMbeans.add(mBeanName.toString());

                }

            } else if (MBeanServerNotification.UNREGISTRATION_NOTIFICATION.equals(mbs.getType())) {

                if (allowedMbeans.contains(mbs.getMBeanName().toString())) {

                    allowedMbeans.remove(mbs.getMBeanName().toString());

                }

            }

        } catch (Exception e) {

        }

    }

    private MBeanServer getMBeanServerProxy(MBeanServer mbeaServer, Subject subject) {

        return (MBeanServer) Proxy.newProxyInstance(MBeanServer.class.getClassLoader(),

                new Class<?>[]{MBeanServer.class},

                new AuthInvocationHandler(mbeaServer, subject));

    }

    private static class HttpUtil {

        public static String getRequestCharset(HttpExchange ex) {

            String charset = null;

            List<String> contentType = ex.getRequestHeaders().get("Content-type");

            if (contentType != null) {

                for (String kv : contentType) {

                    for (String value : kv.split(";")) {

                        value = value.trim();

                        if (value.toLowerCase().startsWith("charset=")) {

                            charset = value.substring("charset=".length());

                        }

                    }

                }

            }

            return charset;

        }

        public static String getAcceptCharset(HttpExchange ex) {

            List<String> acceptCharset = ex.getRequestHeaders().get("Accept-Charset");

            if (acceptCharset != null && acceptCharset.size() > 0) {

                return acceptCharset.get(0);

            }

            return null;

        }

        public static String getGetRequestResource(HttpExchange ex, String charset) throws UnsupportedEncodingException {

            String httpHandlerPath = ex.getHttpContext().getPath();

            String requestURIpath = ex.getRequestURI().getPath();

            String getRequestPath = requestURIpath.substring(httpHandlerPath.length());

            if (charset != null) {

                return URLDecoder.decode(getRequestPath, charset);

            } else {

                return getRequestPath;

            }

        }

        public static String getGetRequestQuery(HttpExchange ex, String charset) throws UnsupportedEncodingException {

            String query = ex.getRequestURI().getQuery();

            if (charset != null && query != null) {

                return URLDecoder.decode(query, charset);

            } else {

                return query;

            }

        }

    }

    private class BasicAuthenticator extends Authenticator {

        @Override

        public Authenticator.Result authenticate(HttpExchange he) {

            Headers rmap = he.getRequestHeaders();

            String auth = rmap.getFirst("Authorization");

            if (auth == null) {

                Headers map = he.getResponseHeaders();

                map.set("WWW-Authenticate", "Basic realm=" + realm);

                return new Authenticator.Retry(401);

            }

            int sp = auth.indexOf(' ');

            if (sp == -1 || !auth.substring(0, sp).equals("Basic")) {

                return new Authenticator.Failure(401);

            }

            byte[] b = Base64.getDecoder().decode(auth.substring(sp + 1));

            String credentials = new String(b);

            int colon = credentials.indexOf(':');

            String uname = credentials.substring(0, colon);

            String pass = credentials.substring(colon + 1);

            if (authMBeanServer.containsKey(credentials)) {

                return new Authenticator.Success(new HttpPrincipal(uname, realm));

            } else {

                Subject subject = null;

                if (authenticator != null) {

                    String[] credential = new String[]{uname, pass};

                    try {

                        subject = authenticator.authenticate(credential);

                    } catch (SecurityException e) {

                        return new Authenticator.Failure(400);

                    }

                }

                MBeanServer proxy = getMBeanServerProxy(mbeanServer, subject);

                authMBeanServer.put(credentials, proxy);

                return new Authenticator.Success(new HttpPrincipal(uname, realm));

            }

        }

    }

    private class AuthInvocationHandler implements InvocationHandler {

        private final MBeanServer mbeanServer;

        private final AccessControlContext acc;

        public AuthInvocationHandler(MBeanServer server, Subject subject) {

            this.mbeanServer = server;

            if (subject == null) {

                this.acc = null;

            } else {

                if (SubjectDelegator.checkRemoveCallerContext(subject)) {

                    acc = JMXSubjectDomainCombiner.getDomainCombinerContext(subject);

                } else {

                    acc = JMXSubjectDomainCombiner.getContext(subject);

                }

            }

        }

        @Override

        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {

            if (acc == null) {

                return method.invoke(mbeanServer, args);

            } else {

                PrivilegedAction<Object> op = () -> {

                    try {

                        return method.invoke(mbeanServer, args);

                    } catch (Exception ex) {

                    }

                    return null;

                };

                return AccessController.doPrivileged(op, acc);

            }

        }

    }

    private class DefaultHTTPHandler implements HttpHandler {

        Map<String, List<String>> allowedMbeans = new HashMap<>();

        DefaultHTTPHandler() {

        }

        @Override

        public void handle(HttpExchange he) throws IOException {

            MBeanServer server = mbeanServer;

            if (env.get("jmx.remote.x.authentication") != null) {

                Headers rmap = he.getRequestHeaders();

                String auth = rmap.getFirst("Authorization");

                int sp = auth.indexOf(' ');

                byte[] b = Base64.getDecoder().decode(auth.substring(sp + 1));

                String authCredentials = new String(b);

                server = authMBeanServer.get(authCredentials);

                if (server == null) {

                    throw new IllegalArgumentException("Invalid HTTP request Headers");

                }

            }

            String charset = HttpUtil.getRequestCharset(he);

            try {

                switch (he.getRequestMethod()) {

                    case "GET":

                        JSONObject resp = getHandler.handle(HttpUtil.getGetRequestResource(he, charset),

                                HttpUtil.getGetRequestQuery(he, charset));

                        sendResponse(he, resp.toJsonString(), ((Long) (((JSONPrimitive) resp.get("status")).getValue())).intValue(), charset);

                        break;

                    case "POST":

                        String requestBody = readRequestBody(he, charset);

                        List<JSONObject> responses = postHandler.handle(requestBody);

                        if (responses.size() == 1) {

                            JSONObject jobj = responses.get(0);

                            sendResponse(he, jobj.toJsonString(), ((Long) (((JSONPrimitive) responses.get(0).get("status")).getValue())).intValue(), charset);

                        } else {

                            int finalCode = HttpURLConnection.HTTP_OK;

                            boolean isHttpOkPresent = responses.stream()

                                    .filter(r -> ((Long) (((JSONPrimitive) responses.get(0).get("status")).getValue())).intValue() == HttpURLConnection.HTTP_OK)

                                    .findFirst().isPresent();

                            if (!isHttpOkPresent) {

                                finalCode = ((Long) (((JSONPrimitive) responses.get(0).get("status")).getValue())).intValue();

                            }

                            JSONArray jarr = new JSONArray();

                            responses.forEach(r -> jarr.add(r));

                            JSONObject jobj = new JSONObject();

                            jobj.put("status", new JSONPrimitive(finalCode));

                            jobj.put("result", jarr);

                            String finalResult = jobj.toJsonString();

                            sendResponse(he, finalResult, finalCode, charset);

                        }

                        break;

                    default:

                        sendResponse(he, "Not supported", HttpURLConnection.HTTP_BAD_METHOD, charset);

                        break;

                }

            } catch (Throwable t) {

                t.printStackTrace();

            }

        }

        private String readRequestBody(HttpExchange he, String charset) throws IOException {

            StringBuilder stringBuilder = new StringBuilder();

            InputStreamReader in = charset != null ? new InputStreamReader(he.getRequestBody(), charset) : new InputStreamReader(he.getRequestBody());

            BufferedReader br = new BufferedReader(in);

            String line;

            while ((line = br.readLine()) != null) {

                String decode = charset != null ? URLDecoder.decode(line, charset) : line;

                stringBuilder.append(decode);

            }

            return stringBuilder.toString();

        }

        private void sendResponse(HttpExchange exchange, String response, int code, String charset) throws IOException {

            String acceptCharset = HttpUtil.getAcceptCharset(exchange);

            if (acceptCharset != null) {

                charset = acceptCharset;

            }

            // Set response headers explicitly

            String msg = charset == null ? response : URLEncoder.encode(response, charset);

            byte[] bytes = msg.getBytes();

            Headers resHeaders = exchange.getResponseHeaders();

            resHeaders.add("Content-Type", "application/json; charset=" + charset);

            exchange.sendResponseHeaders(code, bytes.length);

            try (OutputStream os = exchange.getResponseBody()) {

                os.write(bytes);

            }

        }

    }

}