Move salt and ID to base session and remove unused imports

This commit is contained in:
Tulir Asokan 2018-03-02 18:39:04 +02:00
parent e1d7cc541f
commit dc2229fdba
3 changed files with 15 additions and 30 deletions

View File

@ -1,13 +1,18 @@
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
import time import time
import platform import platform
import struct
import os
class Session(ABC): class Session(ABC):
def __init__(self): def __init__(self):
self.id = struct.unpack('q', os.urandom(8))[0]
self._sequence = 0 self._sequence = 0
self._last_msg_id = 0 self._last_msg_id = 0
self._time_offset = 0 self._time_offset = 0
self._salt = 0
system = platform.uname() system = platform.uname()
self._device_model = system.system or 'Unknown' self._device_model = system.system or 'Unknown'
@ -53,16 +58,6 @@ class Session(ABC):
def auth_key(self, value): def auth_key(self, value):
raise NotImplementedError raise NotImplementedError
@property
@abstractmethod
def salt(self):
raise NotImplementedError
@salt.setter
@abstractmethod
def salt(self, value):
raise NotImplementedError
@abstractmethod @abstractmethod
def close(self): def close(self):
raise NotImplementedError raise NotImplementedError
@ -96,6 +91,14 @@ class Session(ABC):
def get_file(self, md5_digest, file_size, cls): def get_file(self, md5_digest, file_size, cls):
raise NotImplementedError raise NotImplementedError
@property
def salt(self):
return self._salt
@salt.setter
def salt(self, value):
self._salt = value
@property @property
def device_model(self): def device_model(self):
return self._device_model return self._device_model

View File

@ -28,10 +28,10 @@ class _SentFileType(Enum):
class MemorySession(Session): class MemorySession(Session):
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self._dc_id = None self._dc_id = None
self._server_address = None self._server_address = None
self._port = None self._port = None
self._salt = None
self._auth_key = None self._auth_key = None
self._files = {} self._files = {}
@ -58,14 +58,6 @@ class MemorySession(Session):
def auth_key(self, value): def auth_key(self, value):
self._auth_key = value self._auth_key = value
@property
def salt(self):
return self._salt
@salt.setter
def salt(self, value):
self._salt = value
def close(self): def close(self):
pass pass

View File

@ -1,21 +1,13 @@
import json import json
import os import os
import platform
import sqlite3 import sqlite3
import struct
import time
from base64 import b64decode from base64 import b64decode
from os.path import isfile as file_exists from os.path import isfile as file_exists
from threading import Lock, RLock from threading import Lock, RLock
from .. import utils
from .abstract import Session
from .memory import MemorySession, _SentFileType from .memory import MemorySession, _SentFileType
from ..crypto import AuthKey from ..crypto import AuthKey
from ..tl import TLObject
from ..tl.types import ( from ..tl.types import (
PeerUser, PeerChat, PeerChannel,
InputPeerUser, InputPeerChat, InputPeerChannel,
InputPhoto, InputDocument InputPhoto, InputDocument
) )
@ -47,8 +39,6 @@ class SQLiteSession(MemorySession):
if not self.filename.endswith(EXTENSION): if not self.filename.endswith(EXTENSION):
self.filename += EXTENSION self.filename += EXTENSION
self.id = struct.unpack('q', os.urandom(8))[0]
# Cross-thread safety # Cross-thread safety
self._seq_no_lock = Lock() self._seq_no_lock = Lock()
self._msg_id_lock = Lock() self._msg_id_lock = Lock()
@ -193,7 +183,7 @@ class SQLiteSession(MemorySession):
self._auth_key = None self._auth_key = None
c.close() c.close()
@Session.auth_key.setter @MemorySession.auth_key.setter
def auth_key(self, value): def auth_key(self, value):
self._auth_key = value self._auth_key = value
self._update_session_table() self._update_session_table()