/*
 * Copyright 2017-2018 Baidu Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.baidu.openrasp.hook.server.jetty;

import com.baidu.openrasp.HookHandler;
import com.baidu.openrasp.cloud.model.ErrorType;
import com.baidu.openrasp.cloud.utils.CloudUtils;
import com.baidu.openrasp.hook.server.ServerXssHook;
import com.baidu.openrasp.plugin.checker.CheckParameter;
import com.baidu.openrasp.tool.Reflection;
import com.baidu.openrasp.tool.annotation.HookAnnotation;
import com.baidu.openrasp.tool.model.ApplicationModel;
import javassist.CannotCompileException;
import javassist.CtClass;
import javassist.NotFoundException;

import java.io.IOException;
import java.util.HashMap;

/**
 * @author anyang
 * @Description: jetty xss检测的hook点
 * @date 2018/8/1315:13
 */
@HookAnnotation
public class JettyXssHook extends ServerXssHook {

    @Override
    public boolean isClassMatched(String className) {
        return "org/eclipse/jetty/server/AbstractHttpConnection".equals(className);
    }

    @Override
    protected void hookMethod(CtClass ctClass) throws IOException, CannotCompileException, NotFoundException {

        String src = getInvokeStaticSrc(JettyXssHook.class, "getJettyOutputBuffer", "_generator", Object.class);
        insertBefore(ctClass, "completeResponse", "()V", src);

    }


    public static void getJettyOutputBuffer(Object object) {
        HashMap<String, Object> params = new HashMap<String, Object>();
        try {
            Object buffer = Reflection.getSuperField(object, "_buffer");
            if (buffer != null) {
                String content = new String(buffer.toString().getBytes(), "utf-8");
                params.put("html_body", content);
            }
        } catch (Exception e) {
            String message = ApplicationModel.getServerName() + " xss detectde failed";
            int errorCode = ErrorType.HOOK_ERROR.getCode();
            HookHandler.LOGGER.warn(CloudUtils.getExceptionObject(message, errorCode), e);
        }
        if (HookHandler.requestCache.get() != null && !params.isEmpty()) {
            HookHandler.doCheck(CheckParameter.Type.XSS, params);
        }
    }
}

@@ -1168,32 +1168,60 @@
 io.IOException;%0A
+import java.nio.ByteBuffer;%0A
 import java.util
@@ -1533,16 +1533,92 @@
 assName)
+ %7C%7C%0A                %22org/eclipse/jetty/server/HttpChannel%22.equals(className)
 ;%0A    %7D%0A
@@ -1762,16 +1762,17 @@
 ring src
+1
  = getIn
@@ -1918,16 +1918,180 @@
 )V%22, src
+1);%0A        String src2 = getInvokeStaticSrc(JettyXssHook.class, %22getJetty9OutputBuffer%22, %22$1%22, ByteBuffer.class);%0A        insertBefore(ctClass, %22write%22, null, src2
 );%0A%0A    
@@ -2933,11 +2933,793 @@
  %7D%0A    %7D
+%0A%0A    public static void getJetty9OutputBuffer(ByteBuffer buffer) %7B%0A        if (buffer != null) %7B%0A            HashMap%3CString, Object%3E params = new HashMap%3CString, Object%3E();%0A            try %7B%0A                String content = new String(buffer.array());%0A                params.put(%22html_body%22, content);%0A            %7D catch (Exception e) %7B%0A                String message = ApplicationModel.getServerName() + %22 xss detectde failed%22;%0A                int errorCode = ErrorType.HOOK_ERROR.getCode();%0A                HookHandler.LOGGER.warn(CloudUtils.getExceptionObject(message, errorCode), e);%0A            %7D%0A            if (HookHandler.requestCache.get() != null && !params.isEmpty()) %7B%0A                HookHandler.doCheck(CheckParameter.Type.XSS, params);%0A            %7D%0A        %7D%0A    %7D
 %0A%7D%0A

add jetty9 xss