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

import com.jscape.inet.ssh.protocol.v2.marshaling.primitive.StringCodec;
import com.jscape.inet.ssh.protocol.v2.messages.Signature;
import com.jscape.inet.ssh.protocol.v2.messages.SshMsgUserAuthRequest;
import com.jscape.inet.ssh.protocol.v2.messages.SshMsgUserAuthRequestPublicKey;
import com.jscape.inet.ssh.protocol.v2.messages.SshMsgUserAuthRequestPublicKeySignature;
import com.jscape.util.io.ByteArrayInputStream;
import com.jscape.util.io.ByteArrayOutputStream;
import com.jscape.util.io.DataCodec;
import com.jscape.util.io.primitive.BooleanCodec;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.PublicKey;

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

/**
 * @author Alex Usun
 */
public class SshMsgUserAuthRequestPublicKeyCodec
        implements SshMsgUserAuthRequestCodec.MethodCodec {

    private final DataCodec<PublicKey> publicKeyCodec;
    private final DataCodec<Signature> signatureCodec;

    public SshMsgUserAuthRequestPublicKeyCodec(DataCodec<PublicKey> publicKeyCodec, DataCodec<Signature> signatureCodec) {
        assertNotNull(publicKeyCodec);
        this.publicKeyCodec = publicKeyCodec;

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

    @Override
    public SshMsgUserAuthRequest read(InputStream in, String username, String serviceName)
            throws IOException {
        boolean signatureRequest = BooleanCodec.readValue(in);

        if (signatureRequest) {
            String publicKeyAlgorithm = StringCodec.readUsAsciiValue(in);
            byte[] publicKeyData = StringCodec.readValue(in);
            PublicKey publicKey = this.publicKeyCodec.read(new ByteArrayInputStream(publicKeyData));
            Signature signature = this.signatureCodec.read(in);

            return new SshMsgUserAuthRequestPublicKeySignature(username, serviceName, publicKeyAlgorithm, publicKey, signature);
        } else {
            String publicKeyAlgorithm = StringCodec.readUsAsciiValue(in);
            byte[] publicKeyData = StringCodec.readValue(in);
            PublicKey publicKey = this.publicKeyCodec.read(new ByteArrayInputStream(publicKeyData));

            return new SshMsgUserAuthRequestPublicKey(username, serviceName, publicKeyAlgorithm, publicKey);
        }
    }

    @Override
    public void write(SshMsgUserAuthRequest value, OutputStream out)
            throws IOException {
        if (value instanceof SshMsgUserAuthRequestPublicKeySignature) {
            SshMsgUserAuthRequestPublicKeySignature message = (SshMsgUserAuthRequestPublicKeySignature) value;
            BooleanCodec.writeValue(true, out);
            StringCodec.writeUsAsciiValue(message.publicKeyAlgorithm, out);

            ByteArrayOutputStream buffer = new ByteArrayOutputStream();
            this.publicKeyCodec.write(message.publicKey, buffer);
            StringCodec.writeValue(buffer.buffer(), 0, buffer.position(), out);

            this.signatureCodec.write(message.signature, out);
        } else {
            SshMsgUserAuthRequestPublicKey message = (SshMsgUserAuthRequestPublicKey) value;
            BooleanCodec.writeValue(false, out);
            StringCodec.writeUsAsciiValue(message.publicKeyAlgorithm, out);

            ByteArrayOutputStream buffer = new ByteArrayOutputStream();
            this.publicKeyCodec.write(message.publicKey, buffer);
            StringCodec.writeValue(buffer.buffer(), 0, buffer.position(), out);
        }
    }
}
