MadelineProto/mtproto.py
2015-03-13 19:21:49 +03:00

279 lines
11 KiB
Python

# -*- coding: utf-8 -*-
"""
Created on Tue Sep 2 19:26:15 2014
@author: Anton Grigoryev
@author: Sammy Pfeiffer
"""
from binascii import crc32 as originalcrc32
def crc32(data):
return originalcrc32(data) & 0xffffffff
from datetime import datetime
from time import time
import io
import os.path
import json
import socket
import struct
def vis(bs):
"""
Function to visualize byte streams. Split into bytes, print to console.
:param bs: BYTE STRING
"""
bs = bytearray(bs)
symbols_in_one_line = 8
n = len(bs) // symbols_in_one_line
for i in range(n):
print(str(i*symbols_in_one_line)+" | "+" ".join(["%02X" % b for b in bs[i*symbols_in_one_line:(i+1)*symbols_in_one_line]])) # for every 8 symbols line
if not len(bs) % symbols_in_one_line == 0:
print(str((i+1)*symbols_in_one_line)+" | "+" ".join(["%02X" % b for b in bs[(i+1)*symbols_in_one_line:]])+"\n") # for last line
class TlConstructor:
def __init__(self, json_dict):
self.id = int(json_dict['id'])
self.type = json_dict['type']
self.predicate = json_dict['predicate']
self.params = []
# case of vector
for param in json_dict['params']:
if param['type'] == "Vector<long>":
param['type'] = "Vector t"
param['subtype'] = "long"
else:
param['subtype'] = None
self.params.append(param)
class TlMethod:
def __init__(self, json_dict):
self.id = int(json_dict['id'])
self.type = json_dict['type']
self.method = json_dict['method']
self.params = json_dict['params']
class TL:
def __init__(self, filename):
with open(filename, 'r') as f:
TL_dict = json.load(f)
# Read constructors
self.constructors = TL_dict['constructors']
self.constructor_id = {}
self.constructor_type = {}
for elem in self.constructors:
z = TlConstructor(elem)
self.constructor_id[z.id] = z
self.constructor_type[z.predicate] = z
self.methods = TL_dict['methods']
self.method_id = {}
self.method_name = {}
for elem in self.methods:
z = TlMethod(elem)
self.method_id[z.id] = z
self.method_name[z.method] = z
## Loading TL_schema (should be placed in the same directory as mtproto.py
tl = TL(os.path.join(os.path.dirname(__file__), "TL_schema.JSON"))
def serialize_obj(bytes_io, type_, **kwargs):
try:
tl_constructor = tl.constructor_type[type_]
except KeyError:
raise Exception("Could not extract type: %s" % type_)
bytes_io.write(struct.pack('<i', tl_constructor.id))
for arg in tl_constructor.params:
serialize_param(bytes_io, type_=arg['type'], value=kwargs[arg['name']])
def serialize_method(bytes_io, type_, **kwargs):
try:
tl_method = tl.method_name[type_]
except KeyError:
raise Exception("Could not extract type: %s" % type_)
bytes_io.write(struct.pack('<i', tl_method.id))
for arg in tl_method.params:
serialize_param(bytes_io, type_=arg['type'], value=kwargs[arg['name']])
def serialize_param(bytes_io, type_, value):
if type_ == "int":
assert isinstance(value, int)
bytes_io.write(struct.pack('<i', value))
elif type_ == "long":
assert isinstance(value, int)
bytes_io.write(struct.pack('<q', value))
elif type_ in ["int128", "int256"]:
assert isinstance(value, bytes)
bytes_io.write(value)
elif type_ == 'string' or 'bytes':
l = len(value)
if l < 254: # short string format
bytes_io.write(struct.pack('<b', l)) # 1 byte of string
#bytes_io.write(int.to_bytes(l, 1, 'little')) # 1 byte of string
bytes_io.write(value) # string
bytes_io.write(b'\x00'*((-l-1) % 4)) # padding bytes
else:
bytes_io.write(b'\xfe') # byte 254
bytes_io.write(struct.pack('<i', l)[:3]) # 3 bytes of string
#bytes_io.write(int.to_bytes(l, 3, 'little')) # 3 bytes of string
bytes_io.write(value) # string
bytes_io.write(b'\x00'*(-l % 4)) # padding bytes
def deserialize(bytes_io, type_=None, subtype=None):
"""
:type bytes_io: io.BytesIO object
"""
assert isinstance(bytes_io, io.BytesIO)
# Built-in bare types
if type_ == 'int': x = struct.unpack('<i', bytes_io.read(4))[0]
elif type_ == '#': x = struct.unpack('<I', bytes_io.read(4))[0]
elif type_ == 'long': x = struct.unpack('<q', bytes_io.read(8))[0]
elif type_ == 'double': x = struct.unpack('<d', bytes_io.read(8))[0]
elif type_ == 'int128': x = bytes_io.read(16)
elif type_ == 'int256': x = bytes_io.read(32)
elif type_ == 'string' or type_ == 'bytes':
l = struct.unpack('<B', bytes_io.read(1))[0]
assert l <= 254 # In general, 0xFF byte is not allowed here
if l == 254:
# We have a long string
long_len = struct.unpack('<I', bytes_io.read(3)+b'\x00')[0]
x = bytes_io.read(long_len)
bytes_io.read(-long_len % 4) # skip padding bytes
else:
# We have a short string
x = bytes_io.read(l)
bytes_io.read(-(l+1) % 4) # skip padding bytes
assert isinstance(x, bytes)
elif type_ == 'vector':
assert subtype is not None
count = struct.unpack('<l', bytes_io.read(4))[0]
print("count is: " + str(count))
x = [deserialize(bytes_io, type_=subtype) for i in range(count)]
else:
# Boxed types
i = struct.unpack('<i', bytes_io.read(4))[0] # read type ID
try:
tl_elem = tl.constructor_id[i]
except KeyError:
raise Exception("Could not extract type: %s" % type_)
base_boxed_types = ["Vector t", "Int", "Long", "Double", "String", "Int128", "Int256"]
if tl_elem.type in base_boxed_types:
x = deserialize(bytes_io, type_=tl_elem.predicate, subtype=subtype)
else: # other types
x = {}
for arg in tl_elem.params:
x[arg['name']] = deserialize(bytes_io, type_=arg['type'], subtype=arg['subtype'])
return x
class Session:
""" Manages TCP Transport. encryption and message frames """
def __init__(self, ip, port):
# creating socket
self.sock = socket.socket()
self.sock.connect((ip, port))
self.auth_key_id = None
self.number = 0
def __del__(self):
# closing socket when session object is deleted
self.sock.close()
@staticmethod
def header_unencrypted(message):
"""
Creating header for the unencrypted message:
:param message: byte string to send
"""
# Basic instructions: https://core.telegram.org/mtproto/description#unencrypted-message
# Message id: https://core.telegram.org/mtproto/description#message-identifier-msg-id
# http://stackoverflow.com/questions/8777753/converting-datetime-date-to-utc-timestamp-in-python
# to make it work in py2 and py3 (py3 has the timestamp() method but py2 doesnt)
curr_timestamp = (datetime.utcfromtimestamp(time()) - datetime(1970, 1, 1)).total_seconds()
msg_id = int(curr_timestamp*2**30)*4
#msg_id = int(datetime.utcnow().timestamp()*2**30)*4
return (b'\x00\x00\x00\x00\x00\x00\x00\x00' +
struct.pack('<Q', msg_id) +
struct.pack('<L', len(message)))
# TCP Transport
# Instructions may be found here: https://core.telegram.org/mtproto#tcp-transport
# If a payload (packet) needs to be transmitted from server to client or from client to server,
# it is encapsulated as follows: 4 length bytes are added at the front (to include the length,
# the sequence number, and CRC32; always divisible by 4) and 4 bytes with the packet sequence number
# within this TCP connection (the first packet sent is numbered 0, the next one 1, etc.),
# and 4 CRC32 bytes at the end (length, sequence number, and payload together).
def send_message(self, message):
"""
Forming the message frame and sending message to server
:param message: byte string to send
"""
data = self.header_unencrypted(message) + message
step1 = struct.pack('<LL', len(data)+12, self.number) + data
step2 = step1 + struct.pack('<L', crc32(step1))
self.sock.send(step2)
self.number += 1
# Sending message visualisation to console
print('>>')
vis(step2)
def recv_message(self):
"""
Reading socket and receiving message from server. Check the CRC32.
"""
packet_length_data = self.sock.recv(4) # reads how many bytes to read
if len(packet_length_data) > 0: # if we have smth. in the socket
packet_length = struct.unpack("<L", packet_length_data)[0]
packet = self.sock.recv(packet_length - 4) # read the rest of bytes from socket
# Received message visualisation to console
print('<<')
vis(packet_length_data+packet)
(x, auth_key_id, message_id, message_length)= struct.unpack("<L8s8sI", packet[0:24])
data = packet[24:24+message_length]
crc = packet[-4:]
# print("crc is: " + str(crc))
# print("type of crc: " + str(type(crc)))
# print("crc.__repr__(): " + crc.__repr__())
# print("struct.unpack('<L', crc): (next line)")
# print(struct.unpack('<L', crc))
# print("crc32(packet_length_data + packet[0:-4]): " + str(crc32(packet_length_data + packet[0:-4])))
# print("crc32(packet_length_data + packet[0:-4]).__repr__(): " + crc32(packet_length_data + packet[0:-4]).__repr__())
# Checking the CRC32 correctness of received data
if crc32(packet_length_data + packet[0:-4]) == struct.unpack('<L', crc)[0]:
return data
else:
raise Exception("CRC32 was not correct!")
else:
raise Exception("Nothing in the socket!")
def method_call(self, method, **kwargs):
z=io.BytesIO()
serialize_method(z, method, **kwargs)
# z.getvalue() on py2.7 returns str, which means bytes
# on py3.4 returns bytes
# bytearray is closer to the same data type to be shared
z_val = bytearray(z.getvalue())
# print("z_val: " + z_val.__repr__())
# print("z_val type: " + str(type(z_val)))
# print("len of z_val: " + str(len(z_val)))
self.send_message(z_val)
server_answer = self.recv_message()
return deserialize(io.BytesIO(server_answer))