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