3
uQh1                 @   s   d dl Z d dlZd dlZddlmZ ddlmZmZ ddlm	Z	 ddl
mZ ddlmZmZmZmZmZ yd dlZdZW n. ek
r Z zdZeeZW Y ddZ[X nX d	Zd
ZG dd deZdS )    N   )types   )MemorySession_SentFileType)utils)AuthKey)
InputPhotoInputDocumentPeerUserPeerChatPeerChannelz.session   c                   s   e Zd ZdZd4 fdd	Zd5 fdd	Zdd Zed	d
 Z fddZ	e
jjdd Ze
jjdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd  Zd!d" Zed#d$ Zd%d& Zd'd( Zd)d* Zd+d, Zd6d.d/Zd0d1 Zd2d3 Z  Z S )7SQLiteSessionaq  This session contains the required information to login into your
       Telegram account. NEVER give the saved session file to anyone, since
       they would gain instant access to all your messages and contacts.

       If you think the session has been compromised, close all the sessions
       through an official Telegram client to revoke the authorization.
    Nc                s0  t d krtt j  d| _d| _|rF|| _| jjtsF|  jt7  _d | _| j	 }|j
d |j r|j
d |j d }|tk r| j|d |j
d |j
dtf | j  |j
d	 |j }|r|\| _| _| _}| _t|d
| _|j  n:| j|ddddd |j
dtf | j  |j  | j  d S )Nz:memory:TzDselect name from sqlite_master where type='table' and name='version'zselect version from versionr   )oldzdelete from versionzinsert into version values (?)zselect * from sessions)dataz%version (version integer primary key)zsessions (
                    dc_id integer primary key,
                    server_address text,
                    port integer,
                    auth_key blob,
                    takeout_id integer
                )zentities (
                    id integer primary key,
                    hash integer not null,
                    username text,
                    phone integer,
                    name text,
                    date integer
                )a  sent_files (
                    md5_digest blob,
                    file_size integer,
                    type integer,
                    id integer,
                    hash integer,
                    primary key(md5_digest, file_size, type)
                )zupdate_state (
                    id integer primary key,
                    pts integer,
                    qts integer,
                    date integer,
                    seq integer
                ))sqlite3sqlite3_errsuper__init__filenamesave_entitiesendswith	EXTENSION_conn_cursorexecutefetchoneCURRENT_VERSION_upgrade_databasesave_dc_id_server_address_port_takeout_idr   	_auth_keyclose_create_table_update_session_table)selfZ
session_idcversiontuple_key)	__class__ </tmp/pip-build-2nz6shyl/telethon/telethon/sessions/sqlite.pyr   !   sJ    





		zSQLiteSession.__init__c                s   t  j|}| j|_|S )N)r   cloner   )r)   Zto_instanceZcloned)r.   r/   r0   r1   r   s    zSQLiteSession.clonec             C   s   | j  }|dkr|d7 }|dkr>|d7 }|jd | j|d |dkrZ|d7 }| j|d |dkrt|d7 }|jd |d	kr|d7 }|jd
 |dkr|d7 }|jd |j  d S )Nr   r   zdrop table sent_fileszsent_files (
                md5_digest blob,
                file_size integer,
                type integer,
                id integer,
                hash integer,
                primary key(md5_digest, file_size, type)
            )   zupdate_state (
                id integer primary key,
                pts integer,
                qts integer,
                date integer,
                seq integer
            )   z2alter table sessions add column takeout_id integer   zdelete from entities   z,alter table entities add column date integer)r   r   r'   r&   )r)   r   r*   r/   r/   r0   r   w   s,    



zSQLiteSession._upgrade_databasec             G   s"   x|D ]}| j dj| qW d S )Nzcreate table {})r   format)r*   definitionsZ
definitionr/   r/   r0   r'      s    
zSQLiteSession._create_tablec                sJ   t  j||| | j  | jd}|r@|d r@t|d d| _nd | _d S )Nzselect auth_key from sessionsr   )r   )r   set_dcr(   _executer   r%   )r)   Zdc_idserver_addressportrow)r.   r/   r0   r8      s    
zSQLiteSession.set_dcc             C   s   || _ | j  d S )N)r%   r(   )r)   valuer/   r/   r0   auth_key   s    zSQLiteSession.auth_keyc             C   s   || _ | j  d S )N)r$   r(   )r)   r=   r/   r/   r0   
takeout_id   s    zSQLiteSession.takeout_idc             C   sJ   | j  }|jd |jd| j| j| j| jr2| jjnd| jf |j  d S )Nzdelete from sessionsz2insert or replace into sessions values (?,?,?,?,?)    )	r   r   r!   r"   r#   r%   r-   r$   r&   )r)   r*   r/   r/   r0   r(      s    

z#SQLiteSession._update_session_tablec             C   sJ   | j d|}|rF|\}}}}tjj|tjjd}tjj||||ddS d S )Nz9select pts, qts, date, seq from update_state where id = ?)tzr   )unread_count)r9   datetimefromtimestamptimezoneutcr   updatesState)r)   	entity_idr<   ptsqtsdateseqr/   r/   r0   get_update_state   s    zSQLiteSession.get_update_statec             C   s$   | j d||j|j|jj |j d S )Nz6insert or replace into update_state values (?,?,?,?,?))r9   rJ   rK   rL   	timestamprM   )r)   rI   stater/   r/   r0   set_update_state   s    
zSQLiteSession.set_update_statec          
   C   s4   | j  }z|jdj }dd |D S |j  X d S )Nz0select id, pts, qts, date, seq from update_statec          	   s   sL   | ]D}|d  t jj|d |d tjj|d tjjd|d d dfV  qdS )r   r   r   r2   )rA   r3   )rJ   rK   rL   rM   rB   N)r   rG   rH   rC   rD   rE   rF   ).0r<   r/   r/   r0   	<genexpr>   s   z2SQLiteSession.get_update_states.<locals>.<genexpr>)r   r   fetchallr&   )r)   r*   rowsr/   r/   r0   get_update_states   s    zSQLiteSession.get_update_statesc             C   s   | j dk	r| j j  dS )z;Saves the current session object as session_user_id.sessionN)r   commit)r)   r/   r/   r0   r       s    
zSQLiteSession.savec             C   s&   | j dkrtj| jdd| _ | j j S )z8Asserts that the connection is open and returns a cursorNF)Zcheck_same_thread)r   r   connectr   cursor)r)   r/   r/   r0   r      s    

zSQLiteSession._cursorc          
   G   s(   | j  }z|j||j S |j  X dS )z
        Gets a cursor, executes `stmt` and closes the cursor,
        fetching one row afterwards and returning its result.
        N)r   r   r   r&   )r)   stmtvaluesr*   r/   r/   r0   r9      s    zSQLiteSession._executec             C   s2   | j dkr.| jdk	r.| jj  | jj  d| _dS )z4Closes the connection unless we're working in-memoryz:memory:N)r   r   rW   r&   )r)   r/   r/   r0   r&     s
    



zSQLiteSession.closec             C   s8   | j dkrdS ytj| j  dS  tk
r2   dS X dS )z Deletes the current session filez:memory:TFN)r   osremoveOSError)r)   r/   r/   r0   delete	  s    
zSQLiteSession.deletec             C   s   dd t jdD S )zvLists all the sessions of the users who have ever connected
           using this client and never logged out
        c             S   s.   g | ]&}|j trtjjtjj|d  qS )r   )r   r   r\   pathsplitextbasename)rR   fr/   r/   r0   
<listcomp>  s   z/SQLiteSession.list_sessions.<locals>.<listcomp>.)r\   listdir)clsr/   r/   r0   list_sessions  s    zSQLiteSession.list_sessionsc          
      sd   | j s
dS | j|}|sdS | j }z0ttj f  fdd|D }|jd| W d|j  X dS )zq
        Processes all the found entities on the given TLObject,
        unless .save_entities is False.
        Nc                s   g | ]}|  qS r/   r/   )rR   r<   )now_tupr/   r0   rd   ,  s    z2SQLiteSession.process_entities.<locals>.<listcomp>z4insert or replace into entities values (?,?,?,?,?,?))r   Z_entities_to_rowsr   inttimeexecutemanyr&   )r)   ZtlorU   r*   r/   )ri   r0   process_entities  s    
zSQLiteSession.process_entitiesc             C   s   | j d|S )Nz-select id, hash from entities where phone = ?)r9   )r)   Zphoner/   r/   r0   get_entity_rows_by_phone2  s    z&SQLiteSession.get_entity_rows_by_phonec             C   s   | j  }zl|jd|fj }|s$d S t|dkr^|jdd d |jddd |d d
 D  |d d	 |d d fS |j  X d S )Nz6select id, hash, date from entities where username = ?r   c             S   s   | d p
dS )Nr   r   r/   )tr/   r/   r0   <lambda>C  s    z;SQLiteSession.get_entity_rows_by_username.<locals>.<lambda>)r-   z0update entities set username = null where id = ?c             S   s   g | ]}|d  fqS )r   r/   )rR   ro   r/   r/   r0   rd   E  s    z=SQLiteSession.get_entity_rows_by_username.<locals>.<listcomp>r   rq   rq   )r   r   rT   lensortrl   r&   )r)   usernamer*   resultsr/   r/   r0   get_entity_rows_by_username6  s    z)SQLiteSession.get_entity_rows_by_usernamec             C   s   | j d|S )Nz,select id, hash from entities where name = ?)r9   )r)   namer/   r/   r0   get_entity_rows_by_nameK  s    z%SQLiteSession.get_entity_rows_by_nameTc             C   sB   |r| j d|S | j dtjt|tjt|tjt|S d S )Nz*select id, hash from entities where id = ?z1select id, hash from entities where id in (?,?,?))r9   r   Zget_peer_idr   r   r   )r)   idexactr/   r/   r0   get_entity_rows_by_idO  s    z#SQLiteSession.get_entity_rows_by_idc             C   s2   | j d||tj|j}|r.||d |d S d S )NzSselect id, hash from sent_files where md5_digest = ? and file_size = ? and type = ?r   r   )r9   r   	from_typer=   )r)   
md5_digest	file_sizerg   r<   r/   r/   r0   get_file]  s
    zSQLiteSession.get_filec             C   sF   t |ttfstdt| | jd||tjt|j|j	|j
 d S )NzCannot cache %s instancez4insert or replace into sent_files values (?,?,?,?,?))
isinstancer
   r	   	TypeErrortyper9   r   r|   r=   ry   Zaccess_hash)r)   r}   r~   instancer/   r/   r0   
cache_fileg  s    zSQLiteSession.cache_file)N)N)T)!__name__
__module____qualname____doc__r   r1   r   staticmethodr'   r8   r   r>   setterr?   r(   rN   rQ   rV   r    r   r9   r&   r_   classmethodrh   rm   rn   rv   rx   r{   r   r   __classcell__r/   r/   )r.   r0   r      s2   Q(	



r   )rC   r\   rk   tlr   Zmemoryr   r    r   Zcryptor   Ztl.typesr	   r
   r   r   r   r   r   ImportErrorer   r   r   r   r/   r/   r/   r0   <module>   s    