package com.tiesheng.encrypt.config;

import cn.hutool.core.io.IoUtil;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.ECKeyUtil;
import cn.hutool.crypto.SmUtil;
import cn.hutool.crypto.asymmetric.KeyType;
import cn.hutool.crypto.asymmetric.SM2;
import cn.hutool.json.JSONUtil;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.MethodParameter;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.RequestBodyAdvice;

import java.io.InputStream;
import java.lang.reflect.Type;
import java.nio.charset.Charset;

/**
 * @author hao
 */
@ControllerAdvice
public class EncryptRequestBodyAdvice implements RequestBodyAdvice {


    @Autowired
    EncryptConfig encryptConfig;


    @Override
    public boolean supports(MethodParameter methodParameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) {
        return true;
    }

    @Override
    public HttpInputMessage beforeBodyRead(HttpInputMessage inputMessage, MethodParameter parameter, Type targetType,
                                           Class<? extends HttpMessageConverter<?>> converterType) {

        if (!encryptConfig.isEnable()) {
            return inputMessage;
        }

        try {
            return new DecryptHttpInputMessage(inputMessage, encryptConfig.getPrivateQ());
        } catch (Exception ignore) {
        }

        return inputMessage;
    }

    @Override
    public Object afterBodyRead(Object body, HttpInputMessage inputMessage, MethodParameter parameter,
                                Type targetType, Class<? extends HttpMessageConverter<?>> converterType) {
        return body;
    }

    @Override
    public Object handleEmptyBody(Object body, HttpInputMessage inputMessage, MethodParameter parameter,
                                  Type targetType, Class<? extends HttpMessageConverter<?>> converterType) {
        return body;
    }


    static class DecryptHttpInputMessage implements HttpInputMessage {

        private HttpHeaders headers;
        private InputStream body;

        DecryptHttpInputMessage(HttpInputMessage inputMessage, String privateQ) throws Exception {
            this.headers = inputMessage.getHeaders();
            String bodyStr = IoUtil.read(inputMessage.getBody(), CharsetUtil.CHARSET_UTF_8);
            String encryptData = JSONUtil.parseObj(bodyStr).getStr("encryptData");
            if (!StrUtil.isEmpty(encryptData)) {
                // 部分语言加密之后缺少04前缀，如果解密失败，可尝试增加04
                ECPrivateKeyParameters privateKeyParameters = ECKeyUtil.toSm2PrivateParams(privateQ);
                ECPublicKeyParameters publicKeyParameters = ECKeyUtil.getPublicParams(privateKeyParameters);
                SM2 sm2 = SmUtil.sm2(privateKeyParameters, publicKeyParameters);
                String decrypt = sm2.decryptStr(encryptData, KeyType.PrivateKey);
                this.body = IoUtil.toStream(decrypt, Charset.defaultCharset());
            } else {
                this.body = IoUtil.toStream(bodyStr, Charset.defaultCharset());
            }
        }

        @Override
        public InputStream getBody() {
            return this.body;
        }

        @Override
        public HttpHeaders getHeaders() {
            return this.headers;
        }
    }

}
