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