mirror of
https://github.com/TeamPGM/pyrogram.git
synced 2024-11-28 00:56:19 +00:00
Don't throw errors if auth key creation fails; try again instead
This commit is contained in:
Dan
parent
87b2c4b1e7
commit
cbcb1c78c4
@ -31,9 +31,6 @@ from .internals import MsgId, DataCenter
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# TODO: When using TCP connection mode, the server may close it at any time, causing the Auth key creation to fail
|
||||
# The above is true when dealing with temporary keys, although for perm keys it didn't happened, yet.
|
||||
|
||||
class Auth:
|
||||
CURRENT_DH_PRIME = int(
|
||||
"C71CAEB9C6B1C9048E6C522F70F13F73980D40238E3E21C14934D037563D930F"
|
||||
@ -79,168 +76,176 @@ class Auth:
|
||||
https://core.telegram.org/mtproto/auth_key
|
||||
https://core.telegram.org/mtproto/samples-auth_key
|
||||
"""
|
||||
log.info("Start creating a new auth key on DC{}".format(self.dc_id))
|
||||
|
||||
self.connection.connect()
|
||||
# The server may close the connection at any time, causing the auth key creation to fail.
|
||||
# If that happens, just try again until it succeed.
|
||||
while True:
|
||||
try:
|
||||
log.info("Start creating a new auth key on DC{}".format(self.dc_id))
|
||||
|
||||
# Step 1; Step 2
|
||||
nonce = int.from_bytes(urandom(16), "little", signed=True)
|
||||
log.debug("Send req_pq: {}".format(nonce))
|
||||
res_pq = self.send(functions.ReqPq(nonce))
|
||||
log.debug("Got ResPq: {}".format(res_pq.server_nonce))
|
||||
self.connection.connect()
|
||||
|
||||
# Step 3
|
||||
pq = int.from_bytes(res_pq.pq, "big")
|
||||
log.debug("Start PQ factorization: {}".format(pq))
|
||||
start = time.time()
|
||||
g = Prime.decompose(pq)
|
||||
p, q = sorted((g, pq // g)) # p < q
|
||||
log.debug("Done PQ factorization ({}s): {} {}".format(round(time.time() - start, 3), p, q))
|
||||
# Step 1; Step 2
|
||||
nonce = int.from_bytes(urandom(16), "little", signed=True)
|
||||
log.debug("Send req_pq: {}".format(nonce))
|
||||
res_pq = self.send(functions.ReqPq(nonce))
|
||||
log.debug("Got ResPq: {}".format(res_pq.server_nonce))
|
||||
|
||||
# Step 4
|
||||
server_nonce = res_pq.server_nonce
|
||||
new_nonce = int.from_bytes(urandom(32), "little", signed=True)
|
||||
# Step 3
|
||||
pq = int.from_bytes(res_pq.pq, "big")
|
||||
log.debug("Start PQ factorization: {}".format(pq))
|
||||
start = time.time()
|
||||
g = Prime.decompose(pq)
|
||||
p, q = sorted((g, pq // g)) # p < q
|
||||
log.debug("Done PQ factorization ({}s): {} {}".format(round(time.time() - start, 3), p, q))
|
||||
|
||||
data = types.PQInnerData(
|
||||
res_pq.pq,
|
||||
int.to_bytes(p, 4, "big"),
|
||||
int.to_bytes(q, 4, "big"),
|
||||
nonce,
|
||||
server_nonce,
|
||||
new_nonce,
|
||||
).write()
|
||||
# Step 4
|
||||
server_nonce = res_pq.server_nonce
|
||||
new_nonce = int.from_bytes(urandom(32), "little", signed=True)
|
||||
|
||||
sha = sha1(data).digest()
|
||||
padding = urandom(- (len(data) + len(sha)) % 255)
|
||||
data_with_hash = sha + data + padding
|
||||
encrypted_data = RSA.encrypt(data_with_hash, res_pq.server_public_key_fingerprints[0])
|
||||
data = types.PQInnerData(
|
||||
res_pq.pq,
|
||||
int.to_bytes(p, 4, "big"),
|
||||
int.to_bytes(q, 4, "big"),
|
||||
nonce,
|
||||
server_nonce,
|
||||
new_nonce,
|
||||
).write()
|
||||
|
||||
log.debug("Done encrypt data with RSA")
|
||||
sha = sha1(data).digest()
|
||||
padding = urandom(- (len(data) + len(sha)) % 255)
|
||||
data_with_hash = sha + data + padding
|
||||
encrypted_data = RSA.encrypt(data_with_hash, res_pq.server_public_key_fingerprints[0])
|
||||
|
||||
# Step 5. TODO: Handle "server_DH_params_fail". Code assumes response is ok
|
||||
log.debug("Send req_DH_params")
|
||||
server_dh_params = self.send(
|
||||
functions.ReqDhParams(
|
||||
nonce,
|
||||
server_nonce,
|
||||
int.to_bytes(p, 4, "big"),
|
||||
int.to_bytes(q, 4, "big"),
|
||||
res_pq.server_public_key_fingerprints[0],
|
||||
encrypted_data
|
||||
)
|
||||
)
|
||||
log.debug("Done encrypt data with RSA")
|
||||
|
||||
encrypted_answer = server_dh_params.encrypted_answer
|
||||
# Step 5. TODO: Handle "server_DH_params_fail". Code assumes response is ok
|
||||
log.debug("Send req_DH_params")
|
||||
server_dh_params = self.send(
|
||||
functions.ReqDhParams(
|
||||
nonce,
|
||||
server_nonce,
|
||||
int.to_bytes(p, 4, "big"),
|
||||
int.to_bytes(q, 4, "big"),
|
||||
res_pq.server_public_key_fingerprints[0],
|
||||
encrypted_data
|
||||
)
|
||||
)
|
||||
|
||||
server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
|
||||
new_nonce = int.to_bytes(new_nonce, 32, "little", signed=True)
|
||||
encrypted_answer = server_dh_params.encrypted_answer
|
||||
|
||||
tmp_aes_key = (
|
||||
sha1(new_nonce + server_nonce).digest()
|
||||
+ sha1(server_nonce + new_nonce).digest()[:12]
|
||||
)
|
||||
server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
|
||||
new_nonce = int.to_bytes(new_nonce, 32, "little", signed=True)
|
||||
|
||||
tmp_aes_iv = (
|
||||
sha1(server_nonce + new_nonce).digest()[12:]
|
||||
+ sha1(new_nonce + new_nonce).digest() + new_nonce[:4]
|
||||
)
|
||||
tmp_aes_key = (
|
||||
sha1(new_nonce + server_nonce).digest()
|
||||
+ sha1(server_nonce + new_nonce).digest()[:12]
|
||||
)
|
||||
|
||||
server_nonce = int.from_bytes(server_nonce, "little", signed=True)
|
||||
tmp_aes_iv = (
|
||||
sha1(server_nonce + new_nonce).digest()[12:]
|
||||
+ sha1(new_nonce + new_nonce).digest() + new_nonce[:4]
|
||||
)
|
||||
|
||||
answer_with_hash = IGE.decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv)
|
||||
answer = answer_with_hash[20:]
|
||||
server_nonce = int.from_bytes(server_nonce, "little", signed=True)
|
||||
|
||||
server_dh_inner_data = Object.read(BytesIO(answer))
|
||||
answer_with_hash = IGE.decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv)
|
||||
answer = answer_with_hash[20:]
|
||||
|
||||
log.debug("Done decrypting answer")
|
||||
server_dh_inner_data = Object.read(BytesIO(answer))
|
||||
|
||||
dh_prime = int.from_bytes(server_dh_inner_data.dh_prime, "big")
|
||||
delta_time = server_dh_inner_data.server_time - time.time()
|
||||
log.debug("Done decrypting answer")
|
||||
|
||||
log.debug("Delta time: {}".format(round(delta_time, 3)))
|
||||
dh_prime = int.from_bytes(server_dh_inner_data.dh_prime, "big")
|
||||
delta_time = server_dh_inner_data.server_time - time.time()
|
||||
|
||||
# Step 6
|
||||
g = server_dh_inner_data.g
|
||||
b = int.from_bytes(urandom(256), "big")
|
||||
g_b = int.to_bytes(pow(g, b, dh_prime), 256, "big")
|
||||
log.debug("Delta time: {}".format(round(delta_time, 3)))
|
||||
|
||||
retry_id = 0
|
||||
# Step 6
|
||||
g = server_dh_inner_data.g
|
||||
b = int.from_bytes(urandom(256), "big")
|
||||
g_b = int.to_bytes(pow(g, b, dh_prime), 256, "big")
|
||||
|
||||
data = types.ClientDhInnerData(
|
||||
nonce,
|
||||
server_nonce,
|
||||
retry_id,
|
||||
g_b
|
||||
).write()
|
||||
retry_id = 0
|
||||
|
||||
sha = sha1(data).digest()
|
||||
padding = urandom(- (len(data) + len(sha)) % 16)
|
||||
data_with_hash = sha + data + padding
|
||||
encrypted_data = IGE.encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv)
|
||||
data = types.ClientDhInnerData(
|
||||
nonce,
|
||||
server_nonce,
|
||||
retry_id,
|
||||
g_b
|
||||
).write()
|
||||
|
||||
log.debug("Send set_client_DH_params")
|
||||
set_client_dh_params_answer = self.send(
|
||||
functions.SetClientDhParams(
|
||||
nonce,
|
||||
server_nonce,
|
||||
encrypted_data
|
||||
)
|
||||
)
|
||||
sha = sha1(data).digest()
|
||||
padding = urandom(- (len(data) + len(sha)) % 16)
|
||||
data_with_hash = sha + data + padding
|
||||
encrypted_data = IGE.encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv)
|
||||
|
||||
# TODO: Handle "auth_key_aux_hash" if the previous step fails
|
||||
log.debug("Send set_client_DH_params")
|
||||
set_client_dh_params_answer = self.send(
|
||||
functions.SetClientDhParams(
|
||||
nonce,
|
||||
server_nonce,
|
||||
encrypted_data
|
||||
)
|
||||
)
|
||||
|
||||
# Step 7; Step 8
|
||||
g_a = int.from_bytes(server_dh_inner_data.g_a, "big")
|
||||
auth_key = int.to_bytes(pow(g_a, b, dh_prime), 256, "big")
|
||||
server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
|
||||
# TODO: Handle "auth_key_aux_hash" if the previous step fails
|
||||
|
||||
# TODO: Handle errors
|
||||
# Step 7; Step 8
|
||||
g_a = int.from_bytes(server_dh_inner_data.g_a, "big")
|
||||
auth_key = int.to_bytes(pow(g_a, b, dh_prime), 256, "big")
|
||||
server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
|
||||
|
||||
#######################
|
||||
# Security checks
|
||||
#######################
|
||||
# TODO: Handle errors
|
||||
|
||||
assert dh_prime == self.CURRENT_DH_PRIME
|
||||
log.debug("DH parameters check: OK")
|
||||
#######################
|
||||
# Security checks
|
||||
#######################
|
||||
|
||||
# https://core.telegram.org/mtproto/security_guidelines#g-a-and-g-b-validation
|
||||
g_b = int.from_bytes(g_b, "big")
|
||||
assert 1 < g < dh_prime - 1
|
||||
assert 1 < g_a < dh_prime - 1
|
||||
assert 1 < g_b < dh_prime - 1
|
||||
assert 2 ** (2048 - 64) < g_a < dh_prime - 2 ** (2048 - 64)
|
||||
assert 2 ** (2048 - 64) < g_b < dh_prime - 2 ** (2048 - 64)
|
||||
log.debug("g_a and g_b validation: OK")
|
||||
assert dh_prime == self.CURRENT_DH_PRIME
|
||||
log.debug("DH parameters check: OK")
|
||||
|
||||
# https://core.telegram.org/mtproto/security_guidelines#checking-sha1-hash-values
|
||||
answer = server_dh_inner_data.write() # Call .write() to remove padding
|
||||
assert answer_with_hash[:20] == sha1(answer).digest()
|
||||
log.debug("SHA1 hash values check: OK")
|
||||
# https://core.telegram.org/mtproto/security_guidelines#g-a-and-g-b-validation
|
||||
g_b = int.from_bytes(g_b, "big")
|
||||
assert 1 < g < dh_prime - 1
|
||||
assert 1 < g_a < dh_prime - 1
|
||||
assert 1 < g_b < dh_prime - 1
|
||||
assert 2 ** (2048 - 64) < g_a < dh_prime - 2 ** (2048 - 64)
|
||||
assert 2 ** (2048 - 64) < g_b < dh_prime - 2 ** (2048 - 64)
|
||||
log.debug("g_a and g_b validation: OK")
|
||||
|
||||
# https://core.telegram.org/mtproto/security_guidelines#checking-nonce-server-nonce-and-new-nonce-fields
|
||||
# 1st message
|
||||
assert nonce == res_pq.nonce
|
||||
# 2nd message
|
||||
server_nonce = int.from_bytes(server_nonce, "little", signed=True)
|
||||
assert nonce == server_dh_params.nonce
|
||||
assert server_nonce == server_dh_params.server_nonce
|
||||
# 3rd message
|
||||
assert nonce == set_client_dh_params_answer.nonce
|
||||
assert server_nonce == set_client_dh_params_answer.server_nonce
|
||||
server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
|
||||
log.debug("Nonce fields check: OK")
|
||||
# https://core.telegram.org/mtproto/security_guidelines#checking-sha1-hash-values
|
||||
answer = server_dh_inner_data.write() # Call .write() to remove padding
|
||||
assert answer_with_hash[:20] == sha1(answer).digest()
|
||||
log.debug("SHA1 hash values check: OK")
|
||||
|
||||
# Step 9
|
||||
server_salt = IGE.xor(new_nonce[:8], server_nonce[:8])
|
||||
# https://core.telegram.org/mtproto/security_guidelines#checking-nonce-server-nonce-and-new-nonce-fields
|
||||
# 1st message
|
||||
assert nonce == res_pq.nonce
|
||||
# 2nd message
|
||||
server_nonce = int.from_bytes(server_nonce, "little", signed=True)
|
||||
assert nonce == server_dh_params.nonce
|
||||
assert server_nonce == server_dh_params.server_nonce
|
||||
# 3rd message
|
||||
assert nonce == set_client_dh_params_answer.nonce
|
||||
assert server_nonce == set_client_dh_params_answer.server_nonce
|
||||
server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
|
||||
log.debug("Nonce fields check: OK")
|
||||
|
||||
log.debug("Server salt: {}".format(int.from_bytes(server_salt, "little")))
|
||||
# Step 9
|
||||
server_salt = IGE.xor(new_nonce[:8], server_nonce[:8])
|
||||
|
||||
log.info(
|
||||
"Done auth key exchange: {}".format(
|
||||
set_client_dh_params_answer.__class__.__name__
|
||||
)
|
||||
)
|
||||
log.debug("Server salt: {}".format(int.from_bytes(server_salt, "little")))
|
||||
|
||||
self.connection.close()
|
||||
|
||||
return auth_key
|
||||
log.info(
|
||||
"Done auth key exchange: {}".format(
|
||||
set_client_dh_params_answer.__class__.__name__
|
||||
)
|
||||
)
|
||||
except: # TODO: Too broad exception clause
|
||||
log.warning("Auth key creation failed. Let's try again.")
|
||||
continue
|
||||
else:
|
||||
return auth_key
|
||||
finally:
|
||||
self.connection.close()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user