TeamPGM/pyrogram
2
0
Fork 0
mirror of https://github.com/TeamPGM/pyrogram.git synced 2024-11-24 07:51:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Don't throw errors if auth key creation fails; try again instead

Browse Source
This commit is contained in:
Dan 2017-12-09 15:26:33 +01:00
parent 87b2c4b1e7
commit cbcb1c78c4
1 changed files with 138 additions and 133 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

271
pyrogram/session/auth.py
View File

@ -31,9 +31,6 @@ from .internals import MsgId, DataCenter
log = logging.getLogger(__name__) 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: class Auth:
CURRENT_DH_PRIME = int( CURRENT_DH_PRIME = int(
"C71CAEB9C6B1C9048E6C522F70F13F73980D40238E3E21C14934D037563D930F" "C71CAEB9C6B1C9048E6C522F70F13F73980D40238E3E21C14934D037563D930F"
@ -79,168 +76,176 @@ class Auth:
https://core.telegram.org/mtproto/auth_key https://core.telegram.org/mtproto/auth_key
https://core.telegram.org/mtproto/samples-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 self.connection.connect()
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 3 # Step 1; Step 2
pq = int.from_bytes(res_pq.pq, "big") nonce = int.from_bytes(urandom(16), "little", signed=True)
log.debug("Start PQ factorization: {}".format(pq)) log.debug("Send req_pq: {}".format(nonce))
start = time.time() res_pq = self.send(functions.ReqPq(nonce))
g = Prime.decompose(pq) log.debug("Got ResPq: {}".format(res_pq.server_nonce))
p, q = sorted((g, pq // g)) # p < q
log.debug("Done PQ factorization ({}s): {} {}".format(round(time.time() - start, 3), p, q))
# Step 4 # Step 3
server_nonce = res_pq.server_nonce pq = int.from_bytes(res_pq.pq, "big")
new_nonce = int.from_bytes(urandom(32), "little", signed=True) 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( # Step 4
res_pq.pq, server_nonce = res_pq.server_nonce
int.to_bytes(p, 4, "big"), new_nonce = int.from_bytes(urandom(32), "little", signed=True)
int.to_bytes(q, 4, "big"),
nonce,
server_nonce,
new_nonce,
).write()
sha = sha1(data).digest() data = types.PQInnerData(
padding = urandom(- (len(data) + len(sha)) % 255) res_pq.pq,
data_with_hash = sha + data + padding int.to_bytes(p, 4, "big"),
encrypted_data = RSA.encrypt(data_with_hash, res_pq.server_public_key_fingerprints[0]) 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("Done encrypt data with RSA")
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
)
)
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) encrypted_answer = server_dh_params.encrypted_answer
new_nonce = int.to_bytes(new_nonce, 32, "little", signed=True)
tmp_aes_key = ( server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
sha1(new_nonce + server_nonce).digest() new_nonce = int.to_bytes(new_nonce, 32, "little", signed=True)
+ sha1(server_nonce + new_nonce).digest()[:12]
)
tmp_aes_iv = ( tmp_aes_key = (
sha1(server_nonce + new_nonce).digest()[12:] sha1(new_nonce + server_nonce).digest()
+ sha1(new_nonce + new_nonce).digest() + new_nonce[:4] + 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) server_nonce = int.from_bytes(server_nonce, "little", signed=True)
answer = answer_with_hash[20:]
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") log.debug("Done decrypting answer")
delta_time = server_dh_inner_data.server_time - time.time()
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 log.debug("Delta time: {}".format(round(delta_time, 3)))
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")
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( retry_id = 0
nonce,
server_nonce,
retry_id,
g_b
).write()
sha = sha1(data).digest() data = types.ClientDhInnerData(
padding = urandom(- (len(data) + len(sha)) % 16) nonce,
data_with_hash = sha + data + padding server_nonce,
encrypted_data = IGE.encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv) retry_id,
g_b
).write()
log.debug("Send set_client_DH_params") sha = sha1(data).digest()
set_client_dh_params_answer = self.send( padding = urandom(- (len(data) + len(sha)) % 16)
functions.SetClientDhParams( data_with_hash = sha + data + padding
nonce, encrypted_data = IGE.encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv)
server_nonce,
encrypted_data
)
)
# 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 # TODO: Handle "auth_key_aux_hash" if the previous step fails
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 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)
####################### # TODO: Handle errors
# Security checks
#######################
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 assert dh_prime == self.CURRENT_DH_PRIME
g_b = int.from_bytes(g_b, "big") log.debug("DH parameters check: OK")
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-sha1-hash-values # https://core.telegram.org/mtproto/security_guidelines#g-a-and-g-b-validation
answer = server_dh_inner_data.write() # Call .write() to remove padding g_b = int.from_bytes(g_b, "big")
assert answer_with_hash[:20] == sha1(answer).digest() assert 1 < g < dh_prime - 1
log.debug("SHA1 hash values check: OK") 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 # https://core.telegram.org/mtproto/security_guidelines#checking-sha1-hash-values
# 1st message answer = server_dh_inner_data.write() # Call .write() to remove padding
assert nonce == res_pq.nonce assert answer_with_hash[:20] == sha1(answer).digest()
# 2nd message log.debug("SHA1 hash values check: OK")
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")
# Step 9 # https://core.telegram.org/mtproto/security_guidelines#checking-nonce-server-nonce-and-new-nonce-fields
server_salt = IGE.xor(new_nonce[:8], server_nonce[:8]) # 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( log.debug("Server salt: {}".format(int.from_bytes(server_salt, "little")))
"Done auth key exchange: {}".format(
set_client_dh_params_answer.__class__.__name__
)
)
self.connection.close() log.info(
"Done auth key exchange: {}".format(
return auth_key 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()