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