diaspy/connection.py

   1 #!/usr/bin/env python
   2
   3 import re
   4 import requests
   5 import json
   6 import warnings
   7
   8
   9 """This module abstracts connection to pod.
  10 """
  11
  12
  13 class LoginError(Exception):
  14     pass
  15
  16
  17 class TokenError(Exception):
  18     pass
  19
  20
  21 class Connection():
  22     """Object representing connection with the pod.
  23     """
  24     _token_regex = re.compile(r'content="(.*?)"\s+name="csrf-token')
  25     _userinfo_regex = re.compile(r'window.current_user_attributes = ({.*})')
  26
  27     def __init__(self, pod, username='', password='', schema='https'):
  28         """
  29         :param pod: The complete url of the diaspora pod to use.
  30         :type pod: str
  31         :param username: The username used to log in.
  32         :type username: str
  33         :param password: The password used to log in.
  34         :type password: str
  35         """
  36         self.pod = pod
  37         self.session = requests.Session()
  38         self.login_data = {}
  39         self.token = ''
  40         try: self._setlogin(username, password)
  41         except requests.exceptions.MissingSchema:
  42             self.pod = '{0}://{1}'.format(schema, self.pod)
  43             warnings.warn('schema was missing')
  44         finally: pass
  45         try: self._setlogin(username, password)
  46         except Exception as e: raise LoginError('cannot create login data (caused by: {0})'.format(e))
  47
  48     def __repr__(self):
  49         """Returns token string.
  50         It will be easier to change backend if programs will just use:
  51             repr(connection)
  52         instead of calling a specified method.
  53         """
  54         return self.get_token()
  55
  56     def get(self, string, headers={}, params={}):
  57         """This method gets data from session.
  58         Performs additional checks if needed.
  59
  60         Example:
  61             To obtain 'foo' from pod one should call `get('foo')`.
  62
  63         :param string: URL to get without the pod's URL and slash eg. 'stream'.
  64         :type string: str
  65         """
  66         return self.session.get('{0}/{1}'.format(self.pod, string), params=params, headers=headers)
  67
  68     def post(self, string, data, headers={}, params={}):
  69         """This method posts data to session.
  70         Performs additional checks if needed.
  71
  72         Example:
  73             To post to 'foo' one should call `post('foo', data={})`.
  74
  75         :param string: URL to post without the pod's URL and slash eg. 'status_messages'.
  76         :type string: str
  77         :param data: Data to post.
  78         :param headers: Headers (optional).
  79         :type headers: dict
  80         :param params: Parameters (optional).
  81         :type params: dict
  82         """
  83         string = '{0}/{1}'.format(self.pod, string)
  84         request = self.session.post(string, data, headers=headers, params=params)
  85         return request
  86
  87     def put(self, string, data=None, headers={}, params={}):
  88         """This method PUTs to session.
  89         """
  90         string = '{0}/{1}'.format(self.pod, string)
  91         if data is not None: request = self.session.put(string, data, headers=headers, params=params)
  92         else: request = self.session.put(string, headers=headers, params=params)
  93         return request
  94
  95     def delete(self, string, data, headers={}):
  96         """This method lets you send delete request to session.
  97         Performs additional checks if needed.
  98
  99         :param string: URL to use.
 100         :type string: str
 101         :param data: Data to use.
 102         :param headers: Headers to use (optional).
 103         :type headers: dict
 104         """
 105         string = '{0}/{1}'.format(self.pod, string)
 106         request = self.session.delete(string, data=data, headers=headers)
 107         return request
 108
 109     def _setlogin(self, username, password):
 110         """This function is used to set data for login.
 111
 112         .. note::
 113             It should be called before _login() function.
 114         """
 115         self.username, self.password = username, password
 116         self.login_data = {'user[username]': self.username,
 117                            'user[password]': self.password,
 118                            'authenticity_token': self._fetchtoken()}
 119
 120     def _login(self):
 121         """Handles actual login request.
 122         Raises LoginError if login failed.
 123         """
 124         request = self.post('users/sign_in',
 125                             data=self.login_data,
 126                             headers={'accept': 'application/json'})
 127         if request.status_code != 201:
 128             raise LoginError('{0}: login failed'.format(request.status_code))
 129
 130     def login(self, username='', password=''):
 131         """This function is used to log in to a pod.
 132         Will raise LoginError if password or username was not specified.
 133         """
 134         if username and password: self._setlogin(username, password)
 135         if not self.username or not self.password: raise LoginError('password or username not specified')
 136         self._login()
 137
 138     def logout(self):
 139         """Logs out from a pod.
 140         When logged out you can't do anything.
 141         """
 142         self.get('users/sign_out')
 143         self.username = ''
 144         self.token = ''
 145         self.password = ''
 146
 147     def podswitch(self, pod):
 148         """Switches pod from current to another one.
 149         """
 150         self.pod = pod
 151         self._login()
 152
 153     def getUserInfo(self):
 154         """This function returns the current user's attributes.
 155
 156         :returns: dict -- json formatted user info.
 157         """
 158         request = self.get('bookmarklet')
 159         try:
 160             userdata = json.loads(self._userinfo_regex.search(request.text).group(1))
 161         except AttributeError:
 162             raise errors.DiaspyError('cannot find user data')
 163         return userdata
 164
 165     def _fetchtoken(self):
 166         """This method tries to get token string needed for authentication on D*.
 167
 168         :returns: token string
 169         """
 170         request = self.get('stream')
 171         token = self._token_regex.search(request.text).group(1)
 172         self.token = token
 173         return token
 174
 175     def get_token(self, fetch=False):
 176         """This function returns a token needed for authentication in most cases.
 177         Each time it is run a _fetchtoken() is called and refreshed token is stored.
 178
 179         It is more safe to use than _fetchtoken().
 180         By setting new you can request new token or decide to get stored one.
 181         If no token is stored new one will be fatched anyway.
 182
 183         :returns: string -- token used to authenticate
 184         """
 185         try:
 186             if fetch: self._fetchtoken()
 187             if not self.token: self._fetchtoken()
 188         except requests.exceptions.ConnectionError as e:
 189             warnings.warn('{0} was cought: reusing old token'.format(e))
 190         finally:
 191             if not self.token: raise TokenError('cannot obtain token and no previous token found for reuse')
 192         return self.token