<commit_before>#!/usr/bin/env python

import json
import os.path
import hmac
from time import time
from datetime import datetime
from hashlib import sha256, md5

from snapy.utils import (encrypt, decrypt, decrypt_story,
                          make_media_id, request, get_auth_token, 
                          make_request_token, get_attestation, 
                          timestamp, STATIC_TOKEN, get_client_auth_token)

MEDIA_IMAGE = 0
MEDIA_VIDEO = 1
MEDIA_VIDEO_NOAUDIO = 2

FRIEND_CONFIRMED = 0
FRIEND_UNCONFIRMED = 1
FRIEND_BLOCKED = 2
PRIVACY_EVERYONE = 0
PRIVACY_FRIENDS = 1


def is_video(data):
    return True if data[0:2] == b'\x00\x00' else False


def is_image(data):
    return True if data[0:2] == b'\xFF\xD8' else False


def is_zip(data):
    return True if data[0:2] == b'PK' else False


def get_file_extension(media_type):
    if media_type in (MEDIA_VIDEO, MEDIA_VIDEO_NOAUDIO):
        return 'mp4'
    if media_type == MEDIA_IMAGE:
        return 'jpg'
    return ''


def get_media_type(data):
    if is_video(data):
        return MEDIA_VIDEO
    if is_image(data):
        return MEDIA_IMAGE
    if is_zip(data):
        return MEDIA_VIDEO
    return None


def _map_keys(snap):
    return {
        u'id': snap.get('id', None),
        u'media_id': snap.get('c_id', None),
        u'media_type': snap.get('m', None),
        u'time': snap.get('t', None),
        u'sender': snap.get('sn', None),
        u'recipient': snap.get('rp', None),
        u'status': snap.get('st', None),
        u'screenshot_count': snap.get('c', None),
        u'sent': snap.get('sts', None),
        u'opened': snap.get('ts', None)
    }

class Snapchat(object):
    """Construct a :class:`Snapchat` object used for communicating
    with the Snapchat API.

    Usage:

        from snapy import Snapchat
        snapchat = Snapchat()
        snapchat.login('username', 'password', 'gmail_addr', 'gmail_passwd')
        ...

    """
    def __init__(self):
        self.username = None
        self.auth_token = None
        self.gmail = None
        self.gpasswd = None
        self.gauth = None
        self.expiry = datetime.fromtimestamp(0)

    def _request(self, endpoint, data=None, params=None, files=None,
                 raise_for_status=True, req_type='post', moreheaders={}):
        return request(endpoint, self.auth_token, data, params, files,
                       raise_for_status, req_type, moreheaders)

    def _get_device_token(self):
        r = self._request('/loq/device_id',params={'gauth': self._get_gauth()})
        return r.json()
        
    def _unset_auth(self):
        self.username = None
        self.auth_token = None

    def _get_gauth(self):
        """This is the proper way to access self.gauth when using it for an
        API request. This first checks to see if the Bearer token is expired,
        renewing it if needed.
        """
        if datetime.now() >= self.expiry:
            gauth_token = get_auth_token(self.gmail, self.gpasswd)
            self.gauth = gauth_token[0]
            self.expiry = gauth_token[1]

        return self.gauth

    def restore_token(self, username, auth_token, gmail, gpasswd):
        """Restore a Snapchat session from an auth_token parameter
        returned in the response of a login request. Useful for when
        Snapchat breaks the login API.

        :param username     Snapchat username
        :param auth_token   Snapchat auth_token
        :param gmail        Gmail address
        :param gpasswd      Gmail password
        """
        self.username = username
        self.auth_token = auth_token
        self.gmail = gmail
        self.gpasswd = <PASSWORD>
        gauth_token = get_auth_token(gmail, gpasswd)
        self.gauth = gauth_token[0]
        self.expiry = gauth_token[1]

    def login(self, username, password, gmail, gpasswd):
        """Login to Snapchat account
        Returns a dict containing user information on successful login, the
        data returned is similar to get_updates.

        :param username Snapchat username
        :param password Snapchat password
        :param gmail    Gmail address
        :param gpasswd  Gmail password
        """
        self.gmail = gmail
        self.gpasswd = gpasswd

        now = str(timestamp())
        req_token = make_request_token(STATIC_TOKEN, now)
        gauth_token = get_auth_token(gmail, gpasswd)
        self.gauth = gauth_token[0]
        self.expiry = gauth_token[1]
        string = username + "|" + password + "|" + now + "|" + req_token
        dtoken = self._get_device_token()
        self._unset_auth()
        attestation = get_attestation(username, password, now)
        r = self._request('/loq/login', {
            'username': username,
            'password': password,
            'height': 1280,
            'width': 720,
            'max_video_height': 640,
            'max_video_width': 480,
            'dsig': hmac.new(str(dtoken['dtoken1v']),string,sha256).hexdigest()[:20],
            'dtoken1i': dtoken['dtoken1i'],
            'ptoken': "ie",
            'attestation': attestation,
            'sflag': 1,
            'application_id': 'com.snapchat.android',
            'req_token': req_token
        }, {
            'now': now, 
            'gauth': self._get_gauth()
        }, None, True, 'post', {
        'X-Snapchat-Client-Auth': get_client_auth_token(username, password, now)['signature']
        })

        result = r.json()

        if 'updates_response' in result:

            if 'auth_token' in result['updates_response']:
                self.auth_token = result['updates_response']['auth_token']

            if 'username' in result['updates_response']:
                self.username = username

        if self.username is None and self.auth_token is None:
            raise Exception(result.get('message', 'unknown error'))

        return result

    def logout(self):
        """Logout of Snapchat account
        Returns true if logout was successful.
        """
        r = self._request('logout', {'username': self.username})
        return len(r.content) == 0

    def get_updates(self, update_timestamp=0):
        """Get user, friend and snap updates
        Returns a dict containing user, friends and snap information.
<commit_msg>Tries to login 3 times before it gives up<commit_after>#!/usr/bin/env python

import json
import os.path
import hmac
from time import time
from datetime import datetime
from hashlib import sha256, md5

from snapy.utils import (encrypt, decrypt, decrypt_story,
                          make_media_id, request, get_auth_token, 
                          make_request_token, get_attestation, 
                          timestamp, STATIC_TOKEN, get_client_auth_token)

MEDIA_IMAGE = 0
MEDIA_VIDEO = 1
MEDIA_VIDEO_NOAUDIO = 2

FRIEND_CONFIRMED = 0
FRIEND_UNCONFIRMED = 1
FRIEND_BLOCKED = 2
PRIVACY_EVERYONE = 0
PRIVACY_FRIENDS = 1


def is_video(data):
    return True if data[0:2] == b'\x00\x00' else False


def is_image(data):
    return True if data[0:2] == b'\xFF\xD8' else False


def is_zip(data):
    return True if data[0:2] == b'PK' else False


def get_file_extension(media_type):
    if media_type in (MEDIA_VIDEO, MEDIA_VIDEO_NOAUDIO):
        return 'mp4'
    if media_type == MEDIA_IMAGE:
        return 'jpg'
    return ''


def get_media_type(data):
    if is_video(data):
        return MEDIA_VIDEO
    if is_image(data):
        return MEDIA_IMAGE
    if is_zip(data):
        return MEDIA_VIDEO
    return None


def _map_keys(snap):
    return {
        u'id': snap.get('id', None),
        u'media_id': snap.get('c_id', None),
        u'media_type': snap.get('m', None),
        u'time': snap.get('t', None),
        u'sender': snap.get('sn', None),
        u'recipient': snap.get('rp', None),
        u'status': snap.get('st', None),
        u'screenshot_count': snap.get('c', None),
        u'sent': snap.get('sts', None),
        u'opened': snap.get('ts', None)
    }

class Snapchat(object):
    """Construct a :class:`Snapchat` object used for communicating
    with the Snapchat API.

    Usage:

        from snapy import Snapchat
        snapchat = Snapchat()
        snapchat.login('username', 'password', 'gmail_addr', 'gmail_passwd')
        ...

    """
    def __init__(self):
        self.username = None
        self.auth_token = None
        self.gmail = None
        self.gpasswd = None
        self.gauth = None
        self.expiry = datetime.fromtimestamp(0)

    def _request(self, endpoint, data=None, params=None, files=None,
                 raise_for_status=True, req_type='post', moreheaders={}):
        return request(endpoint, self.auth_token, data, params, files,
                       raise_for_status, req_type, moreheaders)

    def _get_device_token(self):
        r = self._request('/loq/device_id',params={'gauth': self._get_gauth()})
        return r.json()
        
    def _unset_auth(self):
        self.username = None
        self.auth_token = None

    def _get_gauth(self):
        """This is the proper way to access self.gauth when using it for an
        API request. This first checks to see if the Bearer token is expired,
        renewing it if needed.
        """
        if datetime.now() >= self.expiry:
            gauth_token = get_auth_token(self.gmail, self.gpasswd)
            self.gauth = gauth_token[0]
            self.expiry = gauth_token[1]

        return self.gauth

    def restore_token(self, username, auth_token, gmail, gpasswd):
        """Restore a Snapchat session from an auth_token parameter
        returned in the response of a login request. Useful for when
        Snapchat breaks the login API.

        :param username     Snapchat username
        :param auth_token   Snapchat auth_token
        :param gmail        Gmail address
        :param gpasswd      Gmail password
        """
        self.username = username
        self.auth_token = auth_token
        self.gmail = gmail
        self.gpasswd = <PASSWORD>
        gauth_token = get_auth_token(gmail, gpasswd)
        self.gauth = gauth_token[0]
        self.expiry = gauth_token[1]

    def login(self, username, password, gmail, gpasswd):
        """Login to Snapchat account
        Returns a dict containing user information on successful login, the
        data returned is similar to get_updates.

        :param username Snapchat username
        :param password <PASSWORD>
        :param gmail    Gmail address
        :param gpasswd  Gmail password
        """
        self.gmail = gmail
        self.gpasswd = <PASSWORD>
        i = 0
        logged_in = False
        while i < 4 and logged_in == False:
            i += 1
            now = str(timestamp())
            req_token = make_request_token(STATIC_TOKEN, now)
            gauth_token = get_auth_token(gmail, gpasswd)
            self.gauth = gauth_token[0]
            self.expiry = gauth_token[1]
            string = username + "|" + password + "|" + now + "|" + req_token
            dtoken = self._get_device_token()
            self._unset_auth()
            attestation = get_attestation(username, password, now)
            r = self._request('/loq/login', {
                'username': username,
                'password': password,
                'height': 1280,
                'width': 720,
                'max_video_height': 640,
                'max_video_width': 480,
                'dsig': hmac.new(str(dtoken['dtoken1v']),string,sha256).hexdigest()[:20],
                'dtoken1i': dtoken['dtoken1i'],
                'ptoken': "ie",
                'attestation': attestation,
                'sflag': 1,
                'application_id': 'com.snapchat.android',
                'req_token': req_token
            }, {
                'now': now, 
                'gauth': self._get_gauth()
            }, None, True, 'post', {
            'X-Snapchat-Client-Auth': get_client_auth_token(username, password, now)['signature']
            })

            result = r.json()

            if 'updates_response' in result:
                logged_in = True
                if 'auth_token' in result['updates_response']:
                    self.auth_token = result['updates_response']['auth_token']

                if 'username' in result['updates_response']:
                    self.username = username

        if self.username is None and self.auth_token is None:
            raise Exception(result.get('message', 'unknown error'))

        return result

    def logout(self):
        """Logout of Snapchat account
        Returns true if logout was successful.
        """
        r = self._request('logout', {'username': self.username})
        return len(r.content) == 0

    def get_updates(self, update_timestamp=0):
        """Get user, friend and snap updates
        Returns a dict containing user, friends and snap information.

{'fraction_non_alphanumeric': 0.10475229072837397, 'fraction_numerical': 0.008852306258735829, 'mean_word_length': 2.1831438457736034, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5855197', 'n_tokens_mistral': 4042, 'n_tokens_neox': 3854, 'n_words': 1009}