Telethon/telethon/network/authenticator.py

200 lines
6.4 KiB
Python

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(req_pq_request.to_bytes())
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 = 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
).to_bytes()
# 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(req_dh_params.to_bytes())
# 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:
print(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 = 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)
).to_bytes()
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(set_client_dh.to_bytes())
# 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)