Telethon/telethon/network/authenticator.py

import os
import time
from hashlib import sha1

from ..tl.types import (
    ResPQ, PQInnerData, ServerDHParamsFail, ServerDHParamsOk,
    ServerDHInnerData, ClientDHInnerData, DhGenOk, DhGenRetry, DhGenFail
)
from .. import helpers as utils
from ..crypto import AES, AuthKey, Factorization
from ..crypto import rsa
from ..errors import SecurityError
from ..extensions import BinaryReader
from ..network import MtProtoPlainSender
from ..tl.functions import (
    ReqPqRequest, ReqDHParamsRequest, SetClientDHParamsRequest
)


def do_authentication(connection, retries=5):
    if not retries or retries < 0:
        retries = 1

    last_error = None
    while retries:
        try:
            return _do_authentication(connection)
        except (SecurityError, AssertionError, NotImplementedError) as e:
            last_error = e
        retries -= 1
    raise last_error


def _do_authentication(connection):
    """Executes the authentication process with the Telegram servers.
       If no error is raised, returns both the authorization key and the
       time offset.
    """
    sender = MtProtoPlainSender(connection)

    # Step 1 sending: PQ Request, endianness doesn't matter since it's random
    req_pq_request = ReqPqRequest(
        nonce=int.from_bytes(os.urandom(16), 'big', signed=True)
    )
    sender.send(bytes(req_pq_request))
    with BinaryReader(sender.receive()) as reader:
        req_pq_request.on_response(reader)

    res_pq = req_pq_request.result
    if not isinstance(res_pq, ResPQ):
        raise AssertionError(res_pq)

    if res_pq.nonce != req_pq_request.nonce:
        raise SecurityError('Invalid nonce from server')

    pq = get_int(res_pq.pq)

    # Step 2 sending: DH Exchange
    p, q = Factorization.factorize(pq)
    p, q = rsa.get_byte_array(min(p, q)), rsa.get_byte_array(max(p, q))
    new_nonce = int.from_bytes(os.urandom(32), 'little', signed=True)

    pq_inner_data = bytes(PQInnerData(
        pq=rsa.get_byte_array(pq), p=p, q=q,
        nonce=res_pq.nonce,
        server_nonce=res_pq.server_nonce,
        new_nonce=new_nonce
    ))

    # sha_digest + data + random_bytes
    cipher_text, target_fingerprint = None, None
    for fingerprint in res_pq.server_public_key_fingerprints:
        cipher_text = rsa.encrypt(fingerprint, pq_inner_data)
        if cipher_text is not None:
            target_fingerprint = fingerprint
            break

    if cipher_text is None:
        raise SecurityError(
            'Could not find a valid key for fingerprints: {}'
            .format(', '.join(
                [str(f) for f in res_pq.server_public_key_fingerprints])
            )
        )

    req_dh_params = ReqDHParamsRequest(
        nonce=res_pq.nonce,
        server_nonce=res_pq.server_nonce,
        p=p, q=q,
        public_key_fingerprint=target_fingerprint,
        encrypted_data=cipher_text
    )
    sender.send(bytes(req_dh_params))

    # Step 2 response: DH Exchange
    with BinaryReader(sender.receive()) as reader:
        req_dh_params.on_response(reader)

    server_dh_params = req_dh_params.result
    if isinstance(server_dh_params, ServerDHParamsFail):
        raise SecurityError('Server DH params fail: TODO')

    if not isinstance(server_dh_params, ServerDHParamsOk):
        raise AssertionError(server_dh_params)

    if server_dh_params.nonce != res_pq.nonce:
        raise SecurityError('Invalid nonce from server')

    if server_dh_params.server_nonce != res_pq.server_nonce:
        raise SecurityError('Invalid server nonce from server')

    # Step 3 sending: Complete DH Exchange
    key, iv = utils.generate_key_data_from_nonce(
        res_pq.server_nonce, new_nonce
    )
    plain_text_answer = AES.decrypt_ige(
        server_dh_params.encrypted_answer, key, iv
    )

    with BinaryReader(plain_text_answer) as reader:
        reader.read(20)  # hash sum
        server_dh_inner = reader.tgread_object()
        if not isinstance(server_dh_inner, ServerDHInnerData):
            raise AssertionError(server_dh_inner)

    if server_dh_inner.nonce != res_pq.nonce:
        raise SecurityError('Invalid nonce in encrypted answer')

    if server_dh_inner.server_nonce != res_pq.server_nonce:
        raise SecurityError('Invalid server nonce in encrypted answer')

    dh_prime = get_int(server_dh_inner.dh_prime, signed=False)
    g_a = get_int(server_dh_inner.g_a, signed=False)
    time_offset = server_dh_inner.server_time - int(time.time())

    b = get_int(os.urandom(256), signed=False)
    gb = pow(server_dh_inner.g, b, dh_prime)
    gab = pow(g_a, b, dh_prime)

    # Prepare client DH Inner Data
    client_dh_inner = bytes(ClientDHInnerData(
        nonce=res_pq.nonce,
        server_nonce=res_pq.server_nonce,
        retry_id=0,  # TODO Actual retry ID
        g_b=rsa.get_byte_array(gb)
    ))

    client_dh_inner_hashed = sha1(client_dh_inner).digest() + client_dh_inner

    # Encryption
    client_dh_encrypted = AES.encrypt_ige(client_dh_inner_hashed, key, iv)

    # Prepare Set client DH params
    set_client_dh = SetClientDHParamsRequest(
        nonce=res_pq.nonce,
        server_nonce=res_pq.server_nonce,
        encrypted_data=client_dh_encrypted,
    )
    sender.send(bytes(set_client_dh))

    # Step 3 response: Complete DH Exchange
    with BinaryReader(sender.receive()) as reader:
        set_client_dh.on_response(reader)

    dh_gen = set_client_dh.result
    if isinstance(dh_gen, DhGenOk):
        if dh_gen.nonce != res_pq.nonce:
            raise SecurityError('Invalid nonce from server')

        if dh_gen.server_nonce != res_pq.server_nonce:
            raise SecurityError('Invalid server nonce from server')

        auth_key = AuthKey(rsa.get_byte_array(gab))
        new_nonce_hash = int.from_bytes(
            auth_key.calc_new_nonce_hash(new_nonce, 1), 'little', signed=True
        )

        if dh_gen.new_nonce_hash1 != new_nonce_hash:
            raise SecurityError('Invalid new nonce hash')

        return auth_key, time_offset

    elif isinstance(dh_gen, DhGenRetry):
        raise NotImplementedError('DhGenRetry')

    elif isinstance(dh_gen, DhGenFail):
        raise NotImplementedError('DhGenFail')

    else:
        raise NotImplementedError('DH Gen unknown: {}'.format(dh_gen))


def get_int(byte_array, signed=True):
    """Gets the specified integer from its byte array.
       This should be used by the authenticator,
       who requires the data to be in big endian
    """
    return int.from_bytes(byte_array, byteorder='big', signed=signed)