diaspy/connection.py

   1 #!/usr/bin/env python
   2
   3 import re
   4 import requests
   5 import json
   6
   7
   8 class LoginError(Exception):
   9     pass
  10
  11
  12 class Connection():
  13     """Object representing connection with the server.
  14     It is pushed around internally and is considered private.
  15     """
  16     _token_regex = re.compile(r'content="(.*?)"\s+name="csrf-token')
  17     _userinfo_regex = re.compile(r'window.current_user_attributes = ({.*})')
  18     login_data = {}
  19
  20     def __init__(self, pod, username='', password=''):
  21         """
  22         :param pod: The complete url of the diaspora pod to use.
  23         :type pod: str
  24         :param username: The username used to log in.
  25         :type username: str
  26         :param password: The password used to log in.
  27         :type password: str
  28         """
  29         self.pod = pod
  30         self.session = requests.Session()
  31         self._setlogin(username, password)
  32
  33     def get(self, string):
  34         """This method gets data from session.
  35         Performs additional checks if needed.
  36
  37         Example:
  38             To obtain 'foo' from pod one should call `get('foo')`.
  39
  40         :param string: URL to get without the pod's URL and slash eg. 'stream'.
  41         :type string: str
  42         """
  43         return self.session.get('{0}/{1}'.format(self.pod, string))
  44
  45     def post(self, string, data, headers={}, params={}):
  46         """This method posts data to session.
  47         Performs additional checks if needed.
  48
  49         Example:
  50             To post to 'foo' one should call `post('foo', data={})`.
  51
  52         :param string: URL to post without the pod's URL and slash eg. 'status_messages'.
  53         :type string: str
  54         :param data: Data to post.
  55         :param headers: Headers (optional).
  56         :type headers: dict
  57         :param params: Parameters (optional).
  58         :type params: dict
  59         """
  60         string = '{0}/{1}'.format(self.pod, string)
  61         request = self.session.post(string, data, headers=headers, params=params)
  62         return request
  63
  64     def delete(self, string, data, headers={}):
  65         """This method lets you send delete request to session.
  66         Performs additional checks if needed.
  67
  68         :param string: URL to use.
  69         :type string: str
  70         :param data: Data to use.
  71         :param headers: Headers to use (optional).
  72         :type headers: dict
  73         """
  74         string = '{0}/{1}'.format(self.pod, string)
  75         request = self.session.delete(string, data=data, headers=headers)
  76         return request
  77
  78     def _setlogin(self, username, password):
  79         """This function is used to set data for login.
  80         .. note::
  81             It should be called before _login() function.
  82         """
  83         self.username, self.password = username, password
  84         self.login_data = {'user[username]': self.username,
  85                            'user[password]': self.password,
  86                            'authenticity_token': self.get_token()}
  87
  88     def _login(self):
  89         """Handles actual login request.
  90         Raises LoginError if login failed.
  91         """
  92         request = self.post('users/sign_in',
  93                             data=self.login_data,
  94                             headers={'accept': 'application/json'})
  95         if request.status_code != 201:
  96             raise LoginError('{0}: Login failed.'.format(request.status_code))
  97
  98     def login(self, username='', password=''):
  99         """This function is used to log in to a pod.
 100         Will raise LoginError if password or username was not specified.
 101         """
 102         if username and password: self._setlogin(username, password)
 103         if not self.username or not self.password: raise LoginError('password or username not specified')
 104         self._login()
 105
 106     def logout(self):
 107         """Logs out from a pod.
 108         When logged out you can't do anything.
 109         """
 110         self.get('users/sign_out')
 111
 112     def podswitch(self, pod):
 113         """Switches pod from current to another one.
 114         """
 115         self.pod = pod
 116         self._login()
 117
 118     def getUserInfo(self):
 119         """This function returns the current user's attributes.
 120
 121         :returns: dict -- json formatted user info.
 122         """
 123         request = self.get('bookmarklet')
 124         userdata = json.loads(self._userinfo_regex.search(request.text).group(1))
 125         return userdata
 126
 127     def get_token(self):
 128         """This function returns a token needed for authentication in most cases.
 129
 130         :returns: string -- token used to authenticate
 131         """
 132         r = self.get('stream')
 133         token = self._token_regex.search(r.text).group(1)
 134         return token