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