package security;

import component.Signature;
import encoding.Block;
import encoding.SelfEncodingBase;
import minsecurity.crypto.AsymKeyException;
import minsecurity.identity.Identity;
import minsecurity.identity.IdentityException;
import org.bouncycastle.crypto.CryptoException;
import packet.CPacket;
import packet.Data;
import packet.Interest;
import packet.MINPacket;

import javax.crypto.BadPaddingException;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.util.Arrays;

/*
 * @Author: hongyu guo
 * @Description: 用于给网络包签名和验签
 * @Version: 1.0.0
 * @Date: 15:49 2021/03/11
 * @Copyright: MIN-Group；国家重大科技基础设施——未来网络北大实验室；深圳市信息论与未来网络重点实验室
 */
public class KeyChain {
    private static final String defaultIdentityName = "/localhost/operator";

    private IdentityManager identityManager;
    private Identity currentIdentity;

    /**
     * 指定当前使用默认的网络身份
     * @param
     * @return
     * @author hongyu guo
     * @date 2021/3/11
    **/
    public KeyChain() throws Exception{
        // identifyManager = new IdentifyManager();
        // currentIdentity = identifyManager.getDefaultIdentity();
        // // DONE: 考虑是否需要在没有默认身份的时候创建一个缺省的本地网络身份
        // if (currentIdentity == null){
        //     Identity newId = this.identifyManager.createIdentityByName(defaultIdentifyName, "", true);
        //     this.identifyManager.setDefaultIdentity(newId);
        //     this.currentIdentity = newId;
        // }
        init(IdentityManager.DefaultIdentityDBPath);
    }

    public KeyChain(String dbPath) throws Exception {
        init(dbPath);
    }

    private void init(String dbPath) throws Exception {
        identityManager = new IdentityManager(dbPath);
        currentIdentity = identityManager.getDefaultIdentity();
        if(identityManager.getDefaultIdentity() == null) {
            Identity defaultIdentity = identityManager.getIdentityByName(defaultIdentityName);
            if(defaultIdentity != null) {
                currentIdentity = defaultIdentity;
                return;
            }

            Identity newIdentity = identityManager.createIdentityByName(defaultIdentityName, "", true);
            if(newIdentity == null) {
                throw new KeyChainException("init failed: can not create identity [" + defaultIdentityName + "]");
            }
            boolean succ = identityManager.setDefaultIdentity(newIdentity, true);
            if(!succ) {
                throw new KeyChainException("init failed: can not set default identity [" + newIdentity.getName() + "]");
            }
            currentIdentity = newIdentity;
        }
    }

    /**
     * 设置当前使用的网络身份，用 password 对目标网络身份进行解锁
     * @param identity
	 * @param passwd 如果passwd不为null 且不为空字符串, 则使用该passwd对identity进行解密
     * @return void
     * @author hongyu guo
     * @date 2021/3/11
    **/
    public void setCurrentIdentity(Identity identity, String passwd) {
        try {
            if(!"".equals(passwd) && identity.isLocked()){
                boolean success = identity.unLock(passwd, identityManager.getPrivateKeyEncryptionAlgorithm());
                if(!success) {
                    throw new KeyChainException("Unlock " + identity.getName() + " by " + passwd + " failed!!");
                }
            }
            currentIdentity = identity;
        } catch (IdentityException | NoSuchAlgorithmException | InvalidKeyException | NoSuchPaddingException
                | BadPaddingException | NoSuchProviderException | IllegalBlockSizeException | KeyChainException e) {
            e.printStackTrace();
        }
    }

    /**
     * 设置默认的的网络身份，不带密码
     * @param identity 身份对象
     * @return boolean
     * @author zhengqi wu
     * @date 2021/6/12
     **/
    public void setDefaultIdentity(Identity identity) throws KeyChainException {
        boolean res = false;
        try {
            res = this.identityManager.setDefaultIdentity(identity, true); // 操作修改到数据库中
        }
        catch (Exception e)
        {
            throw new KeyChainException("Set Default Identity Failed! " + e.getMessage());
        }

        if (!res)
        {
            throw new KeyChainException("Maybe the identity is null.");
        }
        this.currentIdentity = identity;
    }

    /**
     * 以身份名称字符串名称获取身份对象，查询不到返回为空
     * @param identityName 身份名称字符串
     * @return Identity 身份对象
     * @author zhengqi wu
     * @date 2021/6/12
     **/
    public Identity getIdentityByName(String identityName) throws KeyChainException {
        Identity res = null;
        try {
            res = this.identityManager.getIdentityByName(identityName);
        }
        catch (Exception e)
        {
            throw new KeyChainException("Get" + identityName + " failed! " + e.getMessage());
        }
        return res;
    }

    /**
     * 为一个网络身份申请证书
     * @param identity
	 * @param force
     * @return void
     * @author hongyu guo
     * @date 2021/3/11
    **/
    public void generateCertificationForIdentity(Identity identity, boolean force) {
        // TODO: 这边应该发起网络通信，向 MIS 请求给这个网络身份签发一个证书，留待 MIR 完成后进行补充
    }

    /**
     * 检查一个网络身份是否可用
     *
     * 1. 首先检查 identity 是否为空；
     * 2. 接着检查 identity 是否包含私钥；
     * 3. 接着检查 identify 是否被锁定
     * @param identity
     * @return void
     * @author hongyu guo
     * @date 2021/3/11
    **/
    public void checkIdentifyCanUseToSign(Identity identity) throws KeyChainException {
        if(identity == null)
            throw new KeyChainException("identity is null");
        if(!identity.hasPrivateKey()){
            throw new KeyChainException("Identify not have Private key, so can't use to sign!");
        }
        if(identity.isLocked()){
            throw new KeyChainException("Identify is locked, so can't use to sign");
        }
    }


    /**
     * 从 MIN 网络包中提取出标识区和只读区的值，用于签名和验签
     * @param packet MINPacket
     * @return byte[]
     * @throws Exception
     */
     private byte[] getIdentifierAndReadOnlyValueFromPacket(MINPacket packet) throws Exception{
         byte[] rawData = null;
         byte[] iBlockValue = null;
         byte[] rBlockValue = null;
         int totalLength = 0;
         // 1. 获取标识区数据
         Block iBlock = new SelfEncodingBase().selfWireEncode(packet.identifierField);
         if (iBlock != null){
             iBlockValue = iBlock.getValue();
             totalLength += iBlockValue.length;
         }
         // 2. 获取只读区数据
         Block rBlock = new SelfEncodingBase().selfWireEncode(packet.readOnlyField);
         if (rBlock != null){
             rBlockValue = rBlock.getValue();
             totalLength += rBlockValue.length;
         }
         // 3. 合并数据
         if (totalLength == 0)
             return rawData;
         rawData = new byte[totalLength];
         if (iBlock != null)
            System.arraycopy(iBlockValue, 0, rawData, 0, iBlockValue.length);
         if (rBlock != null)
            System.arraycopy(rBlockValue, 0, rawData, iBlockValue.length, rBlockValue.length);
        return rawData;
    }

    /**
     * 给一个通用的网络包签名
     * @param packet
     * @throws Exception
     */
    public void sign(MINPacket packet) throws Exception{
        // 首先检查当前使用的身份是否可以用来签名
        checkIdentifyCanUseToSign(this.currentIdentity);

        // 1. 提取标识区和只读区的值，对其进行签名
        byte[] rawData = getIdentifierAndReadOnlyValueFromPacket(packet);

        // 2. 对标识区和只读区进行签名
        byte[] signResult = this.currentIdentity.sign(rawData);

        // 3. 往包签名区中添加签名
        Signature signature = new Signature(this.currentIdentity, signResult);
        packet.signatureField.addSignature(signature);
    }

    /**
     * 对CPacket进行签名
     * @param cPacket
     * @throws Exception
     */
    public void signCPacket(CPacket cPacket) throws Exception{
        cPacket.fillDataToFields();
        this.sign(cPacket.minPacket);
    }

    /**
     * 对Interest进行签名
     * @param interest
     * @throws Exception
     */
    public void signInterest(Interest interest) throws Exception{
        interest.fillDataToFields();
        this.sign(interest.minPacket);
    }

    /**
     * 对Data进行签名
     * @param data
     * @throws Exception
     */
    public void signData(Data data) throws Exception{
        data.fillDataToFields();
        this.sign(data.minPacket);
    }

    /**
     * 对已经合并好的标识区和只读区进行签名
     * @param mergeArray 原文：已经合并好的标识区和只读区
     * @return 密文
     */
    public byte[] signBytes(byte[] mergeArray) throws Exception {
        return this.currentIdentity.sign(mergeArray);
    }

    /**
     * 验证签名
     * @param mergeArray 原文
     * @param digest 密文
     * @return
     */
    public boolean verifyBytes(byte[] mergeArray,byte[] digest) throws Exception {
        return this.currentIdentity.verify(mergeArray,digest);
    }

    /**
     * 验证一个MIN网络包中的签名是否有效
     * @param minPacket
     * @throws Exception
     */
    public boolean verify(MINPacket minPacket) throws Exception{
        // 提取签名区的第一个签名进行验证（认为签名区的第一个签名为包的签名，包含标识区和只读区签名）
        Signature signature = minPacket.signatureField.getSignature(0);
        // TODO 使用临时方法，待更新
        String identityName = signature.getSigInfo().getKeyLocator().getIdentifier().toUriTemp();
        Identity identity = this.identityManager.getIdentityByName(identityName);

        if (identity == null){
            throw new KeyChainException("Verify failed, could not find the identity");
        }

        byte[] rawData = getIdentifierAndReadOnlyValueFromPacket(minPacket);
        return identity.verify(rawData, signature.getSigValue().getValue());
    }

    /**
     * 验证一个 CPacket 中的签名是否有效
     * @param packet
     * @throws Exception
     */
    public boolean verifyCPacket(CPacket packet) throws Exception{
        packet.fillDataToFields();
        return verify(packet.minPacket);
    }

    /**
     * 验证一个 Interest 中的签名是否有效
     * @param interest
     * @throws Exception
     */
    public boolean verifyInterest(Interest interest) throws Exception{
        interest.fillDataToFields();
        return verify(interest.minPacket);
    }

    /**
     * 验证一个 Data 中的签名是否有效
     * @param data
     * @throws Exception
     */
    public boolean verifyData(Data data) throws Exception{
        data.fillDataToFields();
        return verify(data.minPacket);
    }

    /**
     * 将一个网络身份导出为一个 SafeBag 对象
     * @param identity
     * @param passwd
     * @return SafeBag
     * @throws Exception
     */
    public SafeBag exportSafeBag(Identity identity, String passwd) throws Exception{
        if (identity == null)
            return null;
        byte[] res = identity.dump(passwd);
        SafeBag safeBag = new SafeBag(res);
        return safeBag;
    }

    /**
     * 默认导出当前身份
     * @param passwd
     * @return
     * @throws Exception
     */
    public SafeBag exportSafeBag(String passwd) throws Exception {
        return exportSafeBag(this.getCurrentIdentity(),passwd);
    }

    /**
     * 从一个 SafeBag 中导入网络身份，保存到本地
     * @param safeBag
     * @param passwd
     * @param force
     * @throws Exception
     */
    public void importSafeBag(SafeBag safeBag, String passwd, boolean force) throws Exception{
        if (safeBag == null)
            throw new KeyChainException("SafeBag is null");
        Identity identity = Identity.load(safeBag.getValue(), passwd);

        // 测试代码
        if(this.identityManager.existIdentity(identity.getName())){
            System.out.println("已存在该身份: "+identity.getName());
            System.out.println("已存在该身份: "+identity);
        }

        if (!this.identityManager.existIdentity(identity.getName()) || force){
            this.identityManager.saveIdentity(identity, force, false);
        }else {
            throw new KeyChainException(String.format("Identify %s is already exists!", identity.getName()));
        }
    }

    public IdentityManager getIdentifyManager() {
        return identityManager;
    }

    public void setIdentifyManager(IdentityManager identityManager) {
        this.identityManager = identityManager;
    }

    public Identity getCurrentIdentity() {
        return currentIdentity;
    }

    // public void setCurrentIdentity(Identity currentIdentity) {
    //     this.currentIdentity = currentIdentity;
    // }
}