package com.ttstd.adbloopback.adb;

import android.util.Base64;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.MessageDigest;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.security.interfaces.RSAPublicKey;

/**
 * ADB 密钥对：负责生成 RSA 密钥、对 AUTH 挑战做 SHA1withRSA 签名，
 * 以及把公钥编码成 adbd 可识别的 adb 公钥格式（base64(PKCS#1 RSAPublicKey) + 注释）。
 */
public class AdbKeyPair {
    private final PrivateKey privateKey;
    private final PublicKey publicKey;

    public AdbKeyPair(PrivateKey privateKey, PublicKey publicKey) {
        this.privateKey = privateKey;
        this.publicKey = publicKey;
    }

    public static AdbKeyPair generate() throws Exception {
        KeyPairGenerator gen = KeyPairGenerator.getInstance("RSA");
        gen.initialize(2048, new SecureRandom());
        KeyPair kp = gen.generateKeyPair();
        return new AdbKeyPair(kp.getPrivate(), kp.getPublic());
    }

    public static AdbKeyPair read(File privateFile, File publicFile) throws Exception {
        byte[] pk = readAllBytes(privateFile);
        byte[] pb = readAllBytes(publicFile);
        KeyFactory kf = KeyFactory.getInstance("RSA");
        PrivateKey priv = kf.generatePrivate(new PKCS8EncodedKeySpec(pk));
        PublicKey pub = kf.generatePublic(new X509EncodedKeySpec(pb));
        return new AdbKeyPair(priv, pub);
    }

    public void write(File privateFile, File publicFile) throws Exception {
        KeyFactory kf = KeyFactory.getInstance("RSA");
        byte[] pk = kf.getKeySpec(privateKey, PKCS8EncodedKeySpec.class).getEncoded();
        byte[] pb = kf.getKeySpec(publicKey, X509EncodedKeySpec.class).getEncoded();
        writeAllBytes(privateFile, pk);
        writeAllBytes(publicFile, pb);
    }

    public PrivateKey getPrivateKey() {
        return privateKey;
    }

    public PublicKey getPublicKey() {
        return publicKey;
    }

    /** 对服务端下发的 AUTH token 做签名：SHA1(token) 再用 RSA PKCS#1 v1.5 签名。 */
    public byte[] sign(byte[] token) throws Exception {
        MessageDigest md = MessageDigest.getInstance("SHA-1");
        byte[] hash = md.digest(token);
        Signature sig = Signature.getInstance("SHA1withRSA");
        sig.initSign(privateKey);
        sig.update(hash);
        return sig.sign();
    }

    /**
     * 生成 adb 格式公钥字符串：base64( DER(PKCS#1 RSAPublicKey) ) + " " + 注释 + "\n"
     * adbd 解析时只取 base64 部分，按 (modulus || exponent) 还原 RSA 公钥。
     */
    public String getAdbPublicKey() throws Exception {
        RSAPublicKey rsa = (RSAPublicKey) publicKey;
        byte[] modulus = toDerInteger(rsa.getModulus());
        byte[] exponent = toDerInteger(rsa.getPublicExponent());
        byte[] pkcs1 = encodeSequence(modulus, exponent);
        String b64 = Base64.encodeToString(pkcs1, Base64.NO_WRAP);
        return b64 + " unknown@unknown\n";
    }

    // ---- DER 编码辅助 ----

    private static byte[] toDerInteger(BigInteger v) {
        return derTlv((byte) 0x02, v.toByteArray());
    }

    private static byte[] encodeSequence(byte[]... parts) throws IOException {
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        for (byte[] p : parts) bos.write(p);
        return derTlv((byte) 0x30, bos.toByteArray());
    }

    private static byte[] derTlv(byte tag, byte[] content) {
        return concat(new byte[]{tag}, derLen(content.length), content);
    }

    private static byte[] derLen(int len) {
        if (len < 0x80) {
            return new byte[]{(byte) len};
        } else if (len < 0x100) {
            return new byte[]{(byte) 0x81, (byte) len};
        } else {
            return new byte[]{(byte) 0x82, (byte) (len >> 8), (byte) len};
        }
    }

    private static byte[] concat(byte[] a, byte[] b) {
        byte[] r = new byte[a.length + b.length];
        System.arraycopy(a, 0, r, 0, a.length);
        System.arraycopy(b, 0, r, a.length, b.length);
        return r;
    }

    private static byte[] concat(byte[] a, byte[] b, byte[] c) {
        return concat(concat(a, b), c);
    }

    private static byte[] readAllBytes(File f) throws IOException {
        try (FileInputStream in = new FileInputStream(f)) {
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            byte[] buf = new byte[4096];
            int n;
            while ((n = in.read(buf)) > 0) bos.write(buf, 0, n);
            return bos.toByteArray();
        }
    }

    private static void writeAllBytes(File f, byte[] data) throws IOException {
        try (FileOutputStream out = new FileOutputStream(f)) {
            out.write(data);
        }
    }
}
