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