Overgram/lib/src/main/java/io/github/lonamiwebs/overgram/crypto/Authenticator.java

package io.github.lonamiwebs.overgram.crypto;

import io.github.lonamiwebs.overgram.network.MTProtoPlainSender;
import io.github.lonamiwebs.overgram.tl.Abstract;
import io.github.lonamiwebs.overgram.tl.Functions;
import io.github.lonamiwebs.overgram.tl.Types;
import io.github.lonamiwebs.overgram.utils.BinaryReader;
import io.github.lonamiwebs.overgram.utils.Utils;
import javafx.util.Pair;

import java.io.IOException;
import java.math.BigInteger;
import java.nio.ByteBuffer;

public class Authenticator {

    public static AuthKey doAuthentication(final MTProtoPlainSender sender) throws IOException, ClassNotFoundException {
        final BigInteger nonce = new BigInteger(Utils.randomBytes(16));
        final Types.ResPQ resPq = sender.send(
                new Functions.ReqPqMulti().nonce(nonce));

        if (!resPq.nonce().equals(nonce)) {
            throw new SecurityException("Step 1 response nonce invalid");
        }

        // Note that Java uses big endian, which is correct here
        final BigInteger pq = new BigInteger(resPq.pq());
        final Pair<BigInteger, BigInteger> pqPair = Utils.factorize(pq);

        // Once again, as big endian
        final byte[] p = pqPair.getKey().toByteArray();
        final byte[] q = pqPair.getValue().toByteArray();

        final BigInteger newNonce = new BigInteger(Utils.randomBytes(32));
        final byte[] pqInnerData = new Types.PQInnerData()
                .pq(resPq.pq())
                .p(p)
                .q(q)
                .nonce(resPq.nonce())
                .serverNonce(resPq.serverNonce())
                .newNonce(newNonce)
                .serializeToBytes();

        byte[] cipherText = null;
        long targetFingerprint = 0;
        for (final long fingerprint : resPq.serverPublicKeyFingerprints()) {
            cipherText = RSA.encrypt(fingerprint, pqInnerData);
            if (cipherText != null) {
                targetFingerprint = fingerprint;
                break;
            }
        }

        if (cipherText == null) {
            throw new SecurityException("Step 2 could not find a known RSA key");
        }

        final Abstract.ServerDHParams abstractDhParams = sender.send(
                new Functions.ReqDHParams()
                        .nonce(resPq.nonce())
                        .serverNonce(resPq.serverNonce())
                        .p(p)
                        .q(q)
                        .publicKeyFingerprint(targetFingerprint)
                        .encryptedData(cipherText)
        );

        if (abstractDhParams instanceof Types.ServerDHParamsFail) {
            // We could also check its nonce etc. but not needed
            throw new SecurityException("Step 2 generation of DH params failed");
        }

        final Types.ServerDHParamsOk dhParams = (Types.ServerDHParamsOk) abstractDhParams;
        if (!dhParams.nonce().equals(resPq.nonce())) {
            throw new SecurityException("Step 2 DH params have invalid nonce");
        }
        if (!dhParams.serverNonce().equals(resPq.serverNonce())) {
            throw new SecurityException("Step 2 DH params have invalid server nonce");
        }

        if (dhParams.encryptedAnswer().length % 16 != 0) {
            throw new SecurityException("Step 3 AES block size mismatch");
        }

        final Pair<byte[], byte[]> keyIv = Utils.generateKeyDataFromNonce(resPq.serverNonce(), newNonce);
        final byte[] key = keyIv.getKey();
        final byte[] iv = keyIv.getValue();
        final byte[] plainText = AES.decryptIge(dhParams.encryptedAnswer(), key, iv);

        final BinaryReader dhReader = new BinaryReader(ByteBuffer.wrap(plainText));
        dhReader.seek(20); // hash sum
        final Types.ServerDHInnerData dhInnerData = (Types.ServerDHInnerData) dhReader.readTl();

        if (!dhInnerData.nonce().equals(resPq.nonce())) {
            throw new SecurityException("Step 3 DH data have invalid nonce");
        }
        if (!dhInnerData.serverNonce().equals(resPq.serverNonce())) {
            throw new SecurityException("Step 3 DH data have invalid server nonce");
        }

        // Big endian once again (and unsigned, so contact with a single zero byte)
        final BigInteger dhPrime = new BigInteger(Utils.concat(new byte[1], dhInnerData.dhPrime()));
        final BigInteger ga = new BigInteger(Utils.concat(new byte[1], dhInnerData.gA()));
        final long timeOffset = dhInnerData.serverTime() - (System.currentTimeMillis() / 1000);

        final byte[] bBytes = Utils.randomBytes(257);
        bBytes[0] = 0;
        final BigInteger b = new BigInteger(bBytes);
        final BigInteger gb = new BigInteger(Integer.toString(dhInnerData.g())).modPow(b, dhPrime);
        final BigInteger gab = ga.modPow(b, dhPrime);

        final byte[] clientDhInnerData = new Types.ClientDHInnerData()
                .nonce(resPq.nonce())
                .serverNonce(resPq.serverNonce())
                .retryId(0)
                .gB(gb.toByteArray())
                .serializeToBytes();

        final byte[] clientDhInnerEncrypted = AES.encryptIge(Utils.concat(
                Utils.sha1digest(clientDhInnerData), clientDhInnerData), key, iv);

        final Abstract.SetClientDHParamsAnswer abstractDh = sender.send(
                new Functions.SetClientDHParams()
                        .nonce(resPq.nonce())
                        .serverNonce(resPq.serverNonce())
                        .encryptedData(clientDhInnerEncrypted)
        );

        if (!(abstractDh instanceof Types.DhGenOk)) {
            // Once again we could check the values for other types but it's invalid anyway
            throw new SecurityException("Step 3 DH response was not OK");
        }

        final Types.DhGenOk dhGen = (Types.DhGenOk) abstractDh;
        if (!dhGen.nonce().equals(resPq.nonce())) {
            throw new SecurityException("Step 3 DH gen data have invalid nonce");
        }
        if (!dhGen.serverNonce().equals(resPq.serverNonce())) {
            throw new SecurityException("Step 3 DH gen data have invalid server nonce");
        }

        final AuthKey authKey = new AuthKey(gab.toByteArray());

        // OK, Retry and Fail have numbers 1, 2 and 3 respectively
        final int nonceNumber = 1;

        final BigInteger newNonceHash = authKey.calcNewNonceHash(newNonce, nonceNumber);
        if (!dhGen.newNonceHash1().equals(newNonceHash)) {
            throw new SecurityException("Step 3 DH gen data have invalid nonce hash");
        }

        return authKey; // TODO And time offset
    }
}