| 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 |