diaspy/connection.py

   1 #!/usr/bin/env python
   2 # -*- coding: utf-8 -*-
   3
   4
   5 """This module abstracts connection to pod.
   6 """
   7
   8
   9 import json
  10 import re
  11 import requests
  12 import warnings
  13
  14 from diaspy import errors
  15
  16
  17 DEBUG = True
  18
  19
  20 class Connection():
  21     """Object representing connection with the pod.
  22     """
  23     _token_regex = re.compile(r'name="csrf-token"\s+content="(.*?)"')
  24     _userinfo_regex = re.compile(r'window.current_user_attributes = ({.*})')
  25     # this is for older version of D*
  26     _token_regex_2 = re.compile(r'content="(.*?)"\s+name="csrf-token')
  27     _userinfo_regex_2 = re.compile(r'gon.user=({.*});gon.preloads')
  28     _verify_SSL = True
  29
  30     def __init__(self, pod, username, password, schema='https'):
  31         """
  32         :param pod: The complete url of the diaspora pod to use.
  33         :type pod: str
  34         :param username: The username used to log in.
  35         :type username: str
  36         :param password: The password used to log in.
  37         :type password: str
  38         """
  39         self.pod = pod
  40         self._session = requests.Session()
  41         self._login_data = {'user[remember_me]': 1, 'utf8': '✓'}
  42         self._userdata = {}
  43         self._token = ''
  44         self._diaspora_session = ''
  45         self._cookies = self._fetchcookies()
  46         self._fetch_token_from = 'stream'
  47         try:
  48             #self._setlogin(username, password)
  49             self._login_data = {'user[username]': username,
  50                                 'user[password]': password,
  51                                 'authenticity_token': self._fetchtoken()}
  52             success = True
  53         except requests.exceptions.MissingSchema:
  54             self.pod = '{0}://{1}'.format(schema, self.pod)
  55             warnings.warn('schema was missing')
  56             success = False
  57         finally:
  58             pass
  59         try:
  60             if not success:
  61                 self._login_data = {'user[username]': username,
  62                                     'user[password]': password,
  63                                     'authenticity_token': self._fetchtoken()}
  64         except Exception as e:
  65             raise errors.LoginError('cannot create login data (caused by: {0})'.format(e))
  66
  67     def _fetchcookies(self):
  68         request = self.get('stream')
  69         return request.cookies
  70
  71     def __repr__(self):
  72         """Returns token string.
  73         It will be easier to change backend if programs will just use:
  74             repr(connection)
  75         instead of calling a specified method.
  76         """
  77         return self._fetchtoken()
  78
  79     def get(self, string, headers={}, params={}, direct=False, **kwargs):
  80         """This method gets data from session.
  81         Performs additional checks if needed.
  82
  83         Example:
  84             To obtain 'foo' from pod one should call `get('foo')`.
  85
  86         :param string: URL to get without the pod's URL and slash eg. 'stream'.
  87         :type string: str
  88         :param direct: if passed as True it will not be expanded
  89         :type direct: bool
  90         """
  91         if not direct: url = '{0}/{1}'.format(self.pod, string)
  92         else: url = string
  93         return self._session.get(url, params=params, headers=headers, verify=self._verify_SSL, **kwargs)
  94
  95     def tokenFrom(self, location):
  96         """Sets location for the *next* fetch of CSRF token.
  97         Intended to be used for oneliners like this one:
  98
  99             connection.tokenFrom('somewhere').delete(...)
 100
 101         where the token comes from "somewhere" instead of the
 102         default stream page.
 103         """
 104         self._fetch_token_from = location
 105         return self
 106
 107     def post(self, string, data, headers={}, params={}, **kwargs):
 108         """This method posts data to session.
 109         Performs additional checks if needed.
 110
 111         Example:
 112             To post to 'foo' one should call `post('foo', data={})`.
 113
 114         :param string: URL to post without the pod's URL and slash eg. 'status_messages'.
 115         :type string: str
 116         :param data: Data to post.
 117         :param headers: Headers (optional).
 118         :type headers: dict
 119         :param params: Parameters (optional).
 120         :type params: dict
 121         """
 122         string = '{0}/{1}'.format(self.pod, string)
 123         if 'X-CSRF-Token' not in headers:
 124             headers['X-CSRF-Token'] = self.get_token()
 125         request = self._session.post(string, data, headers=headers, params=params, verify=self._verify_SSL, **kwargs)
 126         return request
 127
 128     def put(self, string, data=None, headers={}, params={}, **kwargs):
 129         """This method PUTs to session.
 130         """
 131         string = '{0}/{1}'.format(self.pod, string)
 132         if 'X-CSRF-Token' not in headers:
 133             headers['X-CSRF-Token'] = self.get_token()
 134         if data is not None: request = self._session.put(string, data, headers=headers, params=params, **kwargs)
 135         else: request = self._session.put(string, headers=headers, params=params, verify=self._verify_SSL, **kwargs)
 136         return request
 137
 138     def delete(self, string, data = None, headers={}, **kwargs):
 139         """This method lets you send delete request to session.
 140         Performs additional checks if needed.
 141
 142         :param string: URL to use.
 143         :type string: str
 144         :param data: Data to use.
 145         :param headers: Headers to use (optional).
 146         :type headers: dict
 147         """
 148         string = '{0}/{1}'.format(self.pod, string)
 149         if 'X-CSRF-Token' not in headers:
 150             headers['X-CSRF-Token'] = self.get_token()
 151         request = self._session.delete(string, data=data, headers=headers, verify=self._verify_SSL, **kwargs)
 152         return request
 153
 154     def _setlogin(self, username, password):
 155         """This function is used to set data for login.
 156
 157         .. note::
 158             It should be called before _login() function.
 159         """
 160         self._login_data = {'user[username]': username,
 161                             'user[password]': password,
 162                             'authenticity_token': self._fetchtoken()}
 163
 164     def _login(self):
 165         """Handles actual login request.
 166         Raises LoginError if login failed.
 167         """
 168         request = self.post('users/sign_in',
 169                             data=self._login_data,
 170                             allow_redirects=False)
 171         if request.status_code != 302:
 172             raise errors.LoginError('{0}: login failed'.format(request.status_code))
 173
 174     def login(self, remember_me=1):
 175         """This function is used to log in to a pod.
 176         Will raise LoginError if password or username was not specified.
 177         """
 178         if not self._login_data['user[username]'] or not self._login_data['user[password]']:
 179             raise errors.LoginError('username and/or password is not specified')
 180         self._login_data['user[remember_me]'] = remember_me
 181         status = self._login()
 182         self._login_data = {}
 183         return self
 184
 185     def logout(self):
 186         """Logs out from a pod.
 187         When logged out you can't do anything.
 188         """
 189         self.get('users/sign_out')
 190         self.token = ''
 191
 192     def podswitch(self, pod, username, password):
 193         """Switches pod from current to another one.
 194         """
 195         self.pod = pod
 196         self._setlogin(username, password)
 197         self._login()
 198
 199     def _fetchtoken(self):
 200         """This method tries to get token string needed for authentication on D*.
 201
 202         :returns: token string
 203         """
 204         request = self.get(self._fetch_token_from)
 205         token = self._token_regex.search(request.text)
 206         if token is None: token = self._token_regex_2.search(request.text)
 207         if token is not None: token = token.group(1)
 208         else: raise errors.TokenError('could not find valid CSRF token')
 209         self._token = token
 210         self._fetch_token_from = 'stream'
 211         return token
 212
 213     def get_token(self, fetch=True):
 214         """This function returns a token needed for authentication in most cases.
 215         **Notice:** using repr() is recommended method for getting token.
 216
 217         Each time it is run a _fetchtoken() is called and refreshed token is stored.
 218
 219         It is more safe to use than _fetchtoken().
 220         By setting new you can request new token or decide to get stored one.
 221         If no token is stored new one will be fetched anyway.
 222
 223         :returns: string -- token used to authenticate
 224         """
 225         try:
 226             if fetch or not self._token: self._fetchtoken()
 227         except requests.exceptions.ConnectionError as e:
 228             warnings.warn('{0} was cought: reusing old token'.format(e))
 229         finally:
 230             if not self._token: raise errors.TokenError('cannot obtain token and no previous token found for reuse')
 231         return self._token
 232
 233     def getSessionToken(self):
 234         """Returns session token string (_diaspora_session).
 235         """
 236         return self._diaspora_session
 237
 238     def getUserData(self):
 239         """Returns user data.
 240         """
 241         request = self.get('bookmarklet')
 242         userdata = self._userinfo_regex.search(request.text)
 243         if userdata is None: userdata = self._userinfo_regex_2.search(request.text)
 244         if userdata is None: raise errors.DiaspyError('cannot find user data')
 245         userdata = userdata.group(1)
 246         return json.loads(userdata)
 247
 248     def set_verify_SSL(self, verify):
 249         """Sets whether there should be an error if a SSL-Certificate could not be verified.
 250         """
 251         self._verify_SSL = verify