def pack(self, message: Message): data = Long(self.current_salt.salt) + self.session_id + message.write() padding = urandom(-(len(data) + 12) % 16 + 12) # 88 = 88 + 0 (outgoing message) msg_key_large = sha256(self.auth_key[88: 88 + 32] + data + padding).digest() msg_key = msg_key_large[8:24] aes_key, aes_iv = KDF(self.auth_key, msg_key, True) return self.auth_key_id + msg_key + AES.ige_encrypt(data + padding, aes_key, aes_iv)
def unpack(self, b: BytesIO) -> Message: assert b.read(8) == self.auth_key_id, b.getvalue() msg_key = b.read(16) aes_key, aes_iv = KDF(self.auth_key, msg_key, False) data = BytesIO(AES.ige_decrypt(b.read(), aes_key, aes_iv)) data.read(8) # https://core.telegram.org/mtproto/security_guidelines#checking-session-id assert data.read(8) == self.session_id message = Message.read(data) # https://core.telegram.org/mtproto/security_guidelines#checking-sha256-hash-value-of-msg-key # https://core.telegram.org/mtproto/security_guidelines#checking-message-length # 96 = 88 + 8 (incoming message) assert msg_key == sha256(self.auth_key[96:96 + 32] + data.getvalue()).digest()[8:24] # https://core.telegram.org/mtproto/security_guidelines#checking-msg-id # TODO: check for lower msg_ids assert message.msg_id % 2 != 0 return message
def create(self): """ https://core.telegram.org/mtproto/auth_key https://core.telegram.org/mtproto/samples-auth_key """ retries_left = self.MAX_RETRIES # The server may close the connection at any time, causing the auth key creation to fail. # If that happens, just try again up to MAX_RETRIES times. while True: self.connection = Connection(self.dc_id, self.test_mode, self.ipv6, self.proxy) try: log.info("Start creating a new auth key on DC{}".format( self.dc_id)) self.connection.connect() # 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.ReqPqMulti(nonce=nonce)) log.debug("Got ResPq: {}".format(res_pq.server_nonce)) log.debug("Server public key fingerprints: {}".format( res_pq.server_public_key_fingerprints)) for i in res_pq.server_public_key_fingerprints: if i in RSA.server_public_keys: log.debug("Using fingerprint: {}".format(i)) public_key_fingerprint = i break else: log.debug("Fingerprint unknown: {}".format(i)) else: raise Exception("Public key not found") # 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 4 server_nonce = res_pq.server_nonce new_nonce = int.from_bytes(urandom(32), "little", signed=True) data = types.PQInnerData( pq=res_pq.pq, p=p.to_bytes(4, "big"), q=q.to_bytes(4, "big"), nonce=nonce, server_nonce=server_nonce, new_nonce=new_nonce, ).write() sha = sha1(data).digest() padding = urandom(-(len(data) + len(sha)) % 255) data_with_hash = sha + data + padding encrypted_data = RSA.encrypt(data_with_hash, public_key_fingerprint) log.debug("Done encrypt data with RSA") # 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=nonce, server_nonce=server_nonce, p=p.to_bytes(4, "big"), q=q.to_bytes(4, "big"), public_key_fingerprint=public_key_fingerprint, encrypted_data=encrypted_data)) encrypted_answer = server_dh_params.encrypted_answer server_nonce = server_nonce.to_bytes(16, "little", signed=True) new_nonce = new_nonce.to_bytes(32, "little", signed=True) tmp_aes_key = (sha1(new_nonce + server_nonce).digest() + sha1(server_nonce + new_nonce).digest()[:12]) tmp_aes_iv = (sha1(server_nonce + new_nonce).digest()[12:] + sha1(new_nonce + new_nonce).digest() + new_nonce[:4]) server_nonce = int.from_bytes(server_nonce, "little", signed=True) answer_with_hash = AES.ige256_decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv) answer = answer_with_hash[20:] server_dh_inner_data = TLObject.read(BytesIO(answer)) log.debug("Done decrypting 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("Delta time: {}".format(round(delta_time, 3))) # Step 6 g = server_dh_inner_data.g b = int.from_bytes(urandom(256), "big") g_b = pow(g, b, dh_prime).to_bytes(256, "big") retry_id = 0 data = types.ClientDHInnerData(nonce=nonce, server_nonce=server_nonce, retry_id=retry_id, g_b=g_b).write() sha = sha1(data).digest() padding = urandom(-(len(data) + len(sha)) % 16) data_with_hash = sha + data + padding encrypted_data = AES.ige256_encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv) log.debug("Send set_client_DH_params") set_client_dh_params_answer = self.send( functions.SetClientDHParams(nonce=nonce, server_nonce=server_nonce, encrypted_data=encrypted_data)) # TODO: Handle "auth_key_aux_hash" if the previous step fails # Step 7; Step 8 g_a = int.from_bytes(server_dh_inner_data.g_a, "big") auth_key = pow(g_a, b, dh_prime).to_bytes(256, "big") server_nonce = server_nonce.to_bytes(16, "little", signed=True) # TODO: Handle errors ####################### # Security checks ####################### assert dh_prime == Prime.CURRENT_DH_PRIME log.debug("DH parameters 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-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#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 = server_nonce.to_bytes(16, "little", signed=True) log.debug("Nonce fields check: OK") # Step 9 server_salt = AES.xor(new_nonce[:8], server_nonce[:8]) log.debug("Server salt: {}".format( int.from_bytes(server_salt, "little"))) log.info("Done auth key exchange: {}".format( set_client_dh_params_answer.__class__.__name__)) except Exception as e: if retries_left: retries_left -= 1 else: raise e time.sleep(1) continue else: return auth_key finally: self.connection.close()
def get_file(self, dc_id: int, id: int = None, access_hash: int = None, volume_id: int = None, local_id: int = None, secret: int = None, version: int = 0, size: int = None, progress: callable = None, progress_args: tuple = None) -> str: with self.media_sessions_lock: session = self.media_sessions.get(dc_id, None) if session is None: if dc_id != self.dc_id: exported_auth = self.send( functions.auth.ExportAuthorization( dc_id=dc_id ) ) session = Session( self, dc_id, Auth(dc_id, self.test_mode, self.ipv6, self._proxy).create(), is_media=True ) session.start() self.media_sessions[dc_id] = session session.send( functions.auth.ImportAuthorization( id=exported_auth.id, bytes=exported_auth.bytes ) ) else: session = Session( self, dc_id, self.auth_key, is_media=True ) session.start() self.media_sessions[dc_id] = session if volume_id: # Photos are accessed by volume_id, local_id, secret location = types.InputFileLocation( volume_id=volume_id, local_id=local_id, secret=secret ) else: # Any other file can be more easily accessed by id and access_hash location = types.InputDocumentFileLocation( id=id, access_hash=access_hash, version=version ) limit = 1024 * 1024 offset = 0 file_name = "" try: r = session.send( functions.upload.GetFile( location=location, offset=offset, limit=limit ) ) if isinstance(r, types.upload.File): with tempfile.NamedTemporaryFile("wb", delete=False) as f: file_name = f.name while True: chunk = r.bytes if not chunk: break f.write(chunk) offset += limit if progress: progress(self, min(offset, size), size, *progress_args) r = session.send( functions.upload.GetFile( location=location, offset=offset, limit=limit ) ) elif isinstance(r, types.upload.FileCdnRedirect): with self.media_sessions_lock: cdn_session = self.media_sessions.get(r.dc_id, None) if cdn_session is None: cdn_session = Session( self, r.dc_id, Auth(r.dc_id, self.test_mode, self.ipv6, self._proxy).create(), is_media=True, is_cdn=True ) cdn_session.start() self.media_sessions[r.dc_id] = cdn_session try: with tempfile.NamedTemporaryFile("wb", delete=False) as f: file_name = f.name while True: r2 = cdn_session.send( functions.upload.GetCdnFile( file_token=r.file_token, offset=offset, limit=limit ) ) if isinstance(r2, types.upload.CdnFileReuploadNeeded): try: session.send( functions.upload.ReuploadCdnFile( file_token=r.file_token, request_token=r2.request_token ) ) except VolumeLocNotFound: break else: continue chunk = r2.bytes # https://core.telegram.org/cdn#decrypting-files decrypted_chunk = AES.ctr256_decrypt( chunk, r.encryption_key, bytearray( r.encryption_iv[:-4] + (offset // 16).to_bytes(4, "big") ) ) hashes = session.send( functions.upload.GetCdnFileHashes( r.file_token, offset ) ) # https://core.telegram.org/cdn#verifying-files for i, h in enumerate(hashes): cdn_chunk = decrypted_chunk[h.limit * i: h.limit * (i + 1)] assert h.hash == sha256(cdn_chunk).digest(), "Invalid CDN hash part {}".format(i) f.write(decrypted_chunk) offset += limit if progress: progress(self, min(offset, size), size, *progress_args) if len(chunk) < limit: break except Exception as e: raise e except Exception as e: log.error(e, exc_info=True) try: os.remove(file_name) except OSError: pass return "" else: return file_name