diaspy.connection.Connection.login() returns Connection() object as well
[diaspy.git] / diaspy / connection.py
index c2d185b0f344fbd3d1abd0b25fd8a4263ff48bd6..50b1d44cf0577f2c05edbbba93904c315c0d1a76 100644 (file)
@@ -1,33 +1,29 @@
 #!/usr/bin/env python
 
-import re
-import requests
-import json
-import warnings
-
 
 """This module abstracts connection to pod.
 """
 
 
-class LoginError(Exception):
-    pass
+import json
+import re
+import requests
+import warnings
 
+from diaspy import errors
 
-class TokenError(Exception):
-    pass
+
+DEBUG = True
 
 
 class Connection():
-    """Object representing connection with the server.
-    It is pushed around internally and is considered private.
+    """Object representing connection with the pod.
     """
     _token_regex = re.compile(r'content="(.*?)"\s+name="csrf-token')
     _userinfo_regex = re.compile(r'window.current_user_attributes = ({.*})')
-    login_data = {}
-    token = ''
+    _userinfo_regex_2 = re.compile(r'gon.user=({.*});gon.preloads')
 
-    def __init__(self, pod, username='', password=''):
+    def __init__(self, pod, username, password, schema='https'):
         """
         :param pod: The complete url of the diaspora pod to use.
         :type pod: str
@@ -37,10 +33,45 @@ class Connection():
         :type password: str
         """
         self.pod = pod
-        self.session = requests.Session()
-        self._setlogin(username, password)
+        self._session = requests.Session()
+        self._login_data = {'user[remember_me]': 1, 'utf8': '✓'}
+        self._userdata = {}
+        self._token = ''
+        self._diaspora_session = ''
+        self._cookies = self._fetchcookies()
+        try:
+            #self._setlogin(username, password)
+            self._login_data = {'user[username]': username,
+                                'user[password]': password,
+                                'authenticity_token': self._fetchtoken()}
+            success = True
+        except requests.exceptions.MissingSchema:
+            self.pod = '{0}://{1}'.format(schema, self.pod)
+            warnings.warn('schema was missing')
+            success = False
+        finally:
+            pass
+        try:
+            if not success:
+                self._login_data = {'user[username]': username,
+                                    'user[password]': password,
+                                    'authenticity_token': self._fetchtoken()}
+        except Exception as e:
+            raise errors.LoginError('cannot create login data (caused by: {0})'.format(e))
+
+    def _fetchcookies(self):
+        request = self.get('stream')
+        return request.cookies
+
+    def __repr__(self):
+        """Returns token string.
+        It will be easier to change backend if programs will just use:
+            repr(connection)
+        instead of calling a specified method.
+        """
+        return self._fetchtoken()
 
-    def get(self, string, headers={}, params={}):
+    def get(self, string, headers={}, params={}, direct=False, **kwargs):
         """This method gets data from session.
         Performs additional checks if needed.
 
@@ -49,10 +80,14 @@ class Connection():
 
         :param string: URL to get without the pod's URL and slash eg. 'stream'.
         :type string: str
+        :param direct: if passed as True it will not be expanded
+        :type direct: bool
         """
-        return self.session.get('{0}/{1}'.format(self.pod, string), params=params, headers=headers)
+        if not direct: url = '{0}/{1}'.format(self.pod, string)
+        else: url = string
+        return self._session.get(url, params=params, headers=headers, **kwargs)
 
-    def post(self, string, data, headers={}, params={}):
+    def post(self, string, data, headers={}, params={}, **kwargs):
         """This method posts data to session.
         Performs additional checks if needed.
 
@@ -68,18 +103,18 @@ class Connection():
         :type params: dict
         """
         string = '{0}/{1}'.format(self.pod, string)
-        request = self.session.post(string, data, headers=headers, params=params)
+        request = self._session.post(string, data, headers=headers, params=params, **kwargs)
         return request
 
-    def put(self, string, data=None, headers={}, params={}):
+    def put(self, string, data=None, headers={}, params={}, **kwargs):
         """This method PUTs to session.
         """
         string = '{0}/{1}'.format(self.pod, string)
-        if data is not None: request = self.session.put(string, data, headers=headers, params=params)
-        else: request = self.session.put(string, headers=headers, params=params)
+        if data is not None: request = self._session.put(string, data, headers=headers, params=params, **kwargs)
+        else: request = self._session.put(string, headers=headers, params=params, **kwargs)
         return request
 
-    def delete(self, string, data, headers={}):
+    def delete(self, string, data, headers={}, **kwargs):
         """This method lets you send delete request to session.
         Performs additional checks if needed.
 
@@ -90,7 +125,7 @@ class Connection():
         :type headers: dict
         """
         string = '{0}/{1}'.format(self.pod, string)
-        request = self.session.delete(string, data=data, headers=headers)
+        request = self._session.delete(string, data=data, headers=headers, **kwargs)
         return request
 
     def _setlogin(self, username, password):
@@ -99,52 +134,56 @@ class Connection():
         .. note::
             It should be called before _login() function.
         """
-        self.username, self.password = username, password
-        self.login_data = {'user[username]': self.username,
-                           'user[password]': self.password,
-                           'authenticity_token': self._fetchtoken()}
+        self._login_data = {'user[username]': username,
+                            'user[password]': password,
+                            'authenticity_token': self._fetchtoken()}
 
     def _login(self):
         """Handles actual login request.
         Raises LoginError if login failed.
         """
         request = self.post('users/sign_in',
-                            data=self.login_data,
-                            headers={'accept': 'application/json'})
-        if request.status_code != 201:
-            raise LoginError('{0}: login failed'.format(request.status_code))
+                            data=self._login_data,
+                            allow_redirects=False)
+        if request.status_code != 302:
+            raise errors.LoginError('{0}: login failed'.format(request.status_code))
 
-    def login(self, username='', password=''):
+    def login(self, remember_me=1):
         """This function is used to log in to a pod.
         Will raise LoginError if password or username was not specified.
         """
-        if username and password: self._setlogin(username, password)
-        if not self.username or not self.password: raise LoginError('password or username not specified')
-        self._login()
+        if not self._login_data['user[username]'] or not self._login_data['user[password]']:
+            raise errors.LoginError('username and/or password is not specified')
+        self._login_data['user[remember_me]'] = remember_me
+        status = self._login()
+        self._login_data = {}
+        return self
 
     def logout(self):
         """Logs out from a pod.
         When logged out you can't do anything.
         """
         self.get('users/sign_out')
-        self.username = ''
         self.token = ''
-        self.password = ''
 
-    def podswitch(self, pod):
+    def podswitch(self, pod, username, password):
         """Switches pod from current to another one.
         """
         self.pod = pod
+        self._setlogin(username, password)
         self._login()
 
-    def getUserInfo(self):
+    def getUserInfo(self, fetch=False):
         """This function returns the current user's attributes.
 
         :returns: dict -- json formatted user info.
         """
         request = self.get('bookmarklet')
-        userdata = json.loads(self._userinfo_regex.search(request.text).group(1))
-        return userdata
+        userdata = self._userinfo_regex.search(request.text)
+        if userdata is None: userdata = self._userinfo_regex_2.search(request.text)
+        if userdata is None: raise errors.DiaspyError('cannot find user data')
+        userdata = userdata.group(1)
+        return json.loads(userdata)
 
     def _fetchtoken(self):
         """This method tries to get token string needed for authentication on D*.
@@ -153,23 +192,30 @@ class Connection():
         """
         request = self.get('stream')
         token = self._token_regex.search(request.text).group(1)
+        self._token = token
         return token
 
-    def get_token(self, new=False):
+    def get_token(self, fetch=True):
         """This function returns a token needed for authentication in most cases.
+        **Notice:** using repr() is recommended method for getting token.
+
         Each time it is run a _fetchtoken() is called and refreshed token is stored.
 
         It is more safe to use than _fetchtoken().
         By setting new you can request new token or decide to get stored one.
-        If no token is stored new one will be fatched anyway.
+        If no token is stored new one will be fetched anyway.
 
         :returns: string -- token used to authenticate
         """
         try:
-            if new: self.token = self._fetchtoken()
-            if not self.token: self.token = self._fetchtoken()
+            if fetch or not self._token: self._fetchtoken()
         except requests.exceptions.ConnectionError as e:
             warnings.warn('{0} was cought: reusing old token'.format(e))
         finally:
-            if not self.token: raise TokenError('cannot obtain token and no previous token found for reuse')
-        return self.token
+            if not self._token: raise errors.TokenError('cannot obtain token and no previous token found for reuse')
+        return self._token
+
+    def getSessionToken(self):
+        """Returns session token string (_diaspora_session).
+        """
+        return self._diaspora_session