package com.jscape.inet.ssh.protocol.v2.marshaling.algorithms;

import com.jscape.crypto.Hash;
import com.jscape.util.ArrayTools;
import com.jscape.util.io.ByteArrayOutputStream;

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

import static com.jscape.util.Parameters.assertNotNull;

/**
 * A key factory for SSH encryption and MAC algorithms.
 *
 * @author Alex Usun
 */
public class KeyFactory {

    private static final byte A = 'A';
    private static final byte B = 'B';
    private static final byte C = 'C';
    private static final byte D = 'D';
    private static final byte E = 'E';
    private static final byte F = 'F';

    private final Hash hash;
    private final byte[] sharedSecret;
    private final byte[] exchangeHash;
    private final byte[] sessionId;

    public KeyFactory(Hash hash, byte[] sharedSecret, byte[] exchangeHash, byte[] sessionId) {
        assertNotNull(hash);
        this.hash = hash;

        assertNotNull(sharedSecret);
        this.sharedSecret = sharedSecret;

        assertNotNull(exchangeHash);
        this.exchangeHash = exchangeHash;

        assertNotNull(sessionId);
        this.sessionId = sessionId;
    }

    public byte[] ivClientToServer(int length) {
        return keyFor(A, length);
    }

    public byte[] ivServerToClient(int length) {
        return keyFor(B, length);
    }

    public byte[] encryptionKeyClientToServer(int length) {
        return keyFor(C, length);
    }

    public byte[] encryptionKeyServerToClient(int length) {
        return keyFor(D, length);
    }

    public byte[] integrityKeyClientToServer(int length) {
        return keyFor(E, length);
    }

    public byte[] integrityKeyServerToClient(int length) {
        return keyFor(F, length);
    }

    @Override
    public String toString() {
        return "KeyFactory {"
                + "hash=" + this.hash
                + ", sharedSecret=" + Arrays.toString(this.sharedSecret)
                + ", exchangeHash=" + Arrays.toString(this.exchangeHash)
                + ", sessionId=" + Arrays.toString(this.sessionId)
                + '}';
    }

    private byte[] keyFor(byte code, int length) {
        byte[] k1 = this.hash.reset().update(this.sharedSecret).update(this.exchangeHash).update(code).update(this.sessionId).compute();
        return k1.length == length ? k1 : k1.length > length ? reduce(k1, length) : enlarge(k1, length);
    }

    private byte[] reduce(byte[] k1, int length) {
        return ArrayTools.copy(k1, new byte[length]);
    }

    private byte[] enlarge(byte[] k1, int length) {
        ByteArrayOutputStream key = new ByteArrayOutputStream();

        try {
            key.write(k1);

            while (key.position() < length) {
                byte[] k2 = this.hash.reset().update(this.sharedSecret).update(this.exchangeHash).update(key.buffer(), 0, key.position()).compute();
                key.write(k2);
            }

            return ArrayTools.copy(key.buffer(), new byte[length]);
        } catch (IOException e) {
            throw new RuntimeException("Unexpected IO error while key enlarging.", e);
        }
    }
}
