/*
 * Copyright 2017-2019 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.file;

import com.baidu.openrasp.HookHandler;
import com.baidu.openrasp.hook.AbstractClassHook;
import com.baidu.openrasp.plugin.checker.CheckParameter;
import com.baidu.openrasp.plugin.js.engine.JSContext;
import com.baidu.openrasp.plugin.js.engine.JSContextFactory;
import com.baidu.openrasp.tool.Reflection;
import com.baidu.openrasp.tool.annotation.HookAnnotation;
import javassist.CannotCompileException;
import javassist.CtClass;
import javassist.NotFoundException;
import org.mozilla.javascript.Scriptable;

import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.util.Arrays;

/**
 * Created by zhuming01 on 5/5/17.
 * All rights reserved
 */
@HookAnnotation
public class DiskFileItemHook extends AbstractClassHook {
    /**
     * (none-javadoc)
     *
     * @see com.baidu.openrasp.hook.AbstractClassHook#getType()
     */
    @Override
    public String getType() {
        return "fileUpload";
    }


    /**
     * (none-javadoc)
     *
     * @see com.baidu.openrasp.hook.AbstractClassHook#isClassMatched(String)
     */
    @Override
    public boolean isClassMatched(String className) {
        return "org/apache/commons/fileupload/disk/DiskFileItem".equals(className);
    }

    @Override
    protected void hookMethod(CtClass ctClass) throws IOException, CannotCompileException, NotFoundException {
        String src = getInvokeStaticSrc(DiskFileItemHook.class, "checkFileUpload",
                "getName(),get(),$1", String.class, byte[].class, Object.class);
        insertAfter(ctClass, "setHeaders", null, src, true);
    }

    /**
     * 文件上传hook点
     *
     * @param name    文件名
     * @param content 文件数据
     */
    public static void checkFileUpload(String name, byte[] content, Object object) {
        if (name != null && content != null && object != null) {
            JSContext cx = JSContextFactory.enterAndInitContext();
            Scriptable params = cx.newObject(cx.getScope());
            params.put("filename", params, name);
            try {
                if (content.length > 4 * 1024) {
                    content = Arrays.copyOf(content, 4 * 1024);
                }
                params.put("content", params, new String(content, "UTF-8"));
            } catch (UnsupportedEncodingException e) {
                e.printStackTrace();
                params.put("content", params, "[rasp error:" + e.getMessage() + "]");
            }
            String customFileName = Reflection.invokeStringMethod(object, "getHeader", new Class[]{String.class}, "content-disposition");
            if (customFileName != null) {
                customFileName = getFileName(customFileName);
            }
            params.put("name", params, customFileName != null ? customFileName : "");

            HookHandler.doCheck(CheckParameter.Type.FILEUPLOAD, params);
        }
    }

    private static String getFileName(String name) {
        String[] fields = name.split(";");
        for (String field : fields) {
            if (!field.contains("filename") && field.contains("name")) {
                int index = field.indexOf("=");
                if (index > 0 && index < field.length() - 1) {
                    return field.substring(index + 1);
                }
            }
        }
        return null;
    }
}

@@ -2092,17 +2092,17 @@
 ,get(),$
-1
+0
 %22, Strin
@@ -2375,22 +2375,20 @@
  Object 
-object
+file
 ) %7B%0A    
@@ -2430,22 +2430,20 @@
 null && 
-object
+file
  != null
@@ -3058,31 +3058,35 @@
     
-String customFileName =
+boolean isFiled = (Boolean)
  Ref
@@ -3103,131 +3103,76 @@
 voke
-String
 Method(
-object, %22getHeader%22, new Class%5B%5D%7BString.class%7D, %22content-disposition%22);%0A            if (customFileName != null
+file, %22isFormField%22, new Class%5B%5D%7B%7D);%0A            if (!isFiled
 ) %7B%0A
@@ -3187,23 +3187,24 @@
         
-customF
+String f
 ileName 
@@ -3209,51 +3209,80 @@
 e = 
-getFileName(customFileName);%0A            %7D%0A
+Reflection.invokeStringMethod(file, %22getFieldName%22, new Class%5B%5D%7B%7D);%0A    
     
@@ -3316,23 +3316,17 @@
 params, 
-customF
+f
 ileName 
@@ -3339,15 +3339,9 @@
 l ? 
-customF
+f
 ileN
@@ -3350,16 +3350,30 @@
 e : %22%22);
+%0A            %7D
 %0A%0A      
@@ -3458,450 +3458,7 @@
    %7D
-%0A%0A    private static String getFileName(String name) %7B%0A        String%5B%5D fields = name.split(%22;%22);%0A        for (String field : fields) %7B%0A            if (!field.contains(%22filename%22) && field.contains(%22name%22)) %7B%0A                int index = field.indexOf(%22=%22);%0A                if (index %3E 0 && index %3C field.length() - 1) %7B%0A                    return field.substring(index + 1);%0A                %7D%0A            %7D%0A        %7D%0A        return null;%0A    %7D
 %0A%7D%0A

modify fileUpload paramter name