Merge pull request capless#70 from capless/armicron_fix_signup

Fixed signup, fixed tests.

README.md

@@ -129,22 +129,39 @@ u = Cognito('your-user-pool-id','your-client-id',
 
 Register a user to the user pool
 
-**Important:** The keyword arguments used for this method depend on your user pool's configuration, and make sure the client id (app id) used has write permissions for the attriubtes you are trying to create. Example, if you want to create a user with a given_name equal to Johnson make sure the client_id you're using has permissions to edit or create given_name for a user in the pool.
+**Important:** The arguments for `add_base_attributes` and `add_custom_attributes` methods depend on your user pool's configuration, and make sure the client id (app id) used has write permissions for the attriubtes you are trying to create. Example, if you want to create a user with a given_name equal to Johnson make sure the client_id you're using has permissions to edit or create given_name for a user in the pool.
 
 
 ```python
 from warrant import Cognito
 
 u = Cognito('your-user-pool-id', 'your-client-id')
 
-u.register('username', 'password', email='[email protected]', some_random_attr='random value')  # **kwargs are the other attributes that should be set ex. email, given_name, family_name
+u.add_base_attributes(email='[email protected]', some_random_attr='random value')
+
+u.register('username', 'password')
+```
+
+Register with custom attributes.
+
+Firstly, add custom attributes on 'General settings -> Attributes' page.
+Secondly, set permissions on 'Generals settings-> App clients-> Show details-> Set attribute read and write permissions' page.
+```python
+from warrant import Cognito
+
+u = Cognito('your-user-pool-id', 'your-client-id')
+
+u.add_base_attributes(email='[email protected]', some_random_attr='random value')
+
+u.add_custom_attributes(state='virginia', city='Centreville')
+
+u.register('username', 'password')
 ```
 ##### Arguments
 
 - **username:** User Pool username
 - **password:** User Pool password
 - **attr_map:** Attribute map to Cognito's attributes
-- **kwargs:** Additional User Pool attributes ex. `**{'email':'[email protected]'}`
 
 
 #### Authenticate ####

warrant/__init__.py

@@ -263,18 +263,17 @@ def check_token(self, renew=True):
         return expired
 
     def add_base_attributes(self, **kwargs):
-        self.base_attributes=kwargs
+        self.base_attributes = kwargs
 
     def add_custom_attributes(self, **kwargs):
-
-        custom_key='custom'
-        custom_attributes={}
+        custom_key = 'custom'
+        custom_attributes = {}
 
-        for old_key,value in kwargs.items():
-            new_key=custom_key+':'+old_key
-            custom_attributes[new_key]= value
+        for old_key, value in kwargs.items():
+            new_key = custom_key + ':' + old_key
+            custom_attributes[new_key] = value
 
-        self.custom_attributes= custom_attributes
+        self.custom_attributes = custom_attributes
 
     def register(self, username, password, attr_map=None):
         """
@@ -286,7 +285,6 @@ def register(self, username, password, attr_map=None):
         :param username: User Pool username
         :param password: User Pool password
         :param attr_map: Attribute map to Cognito's attributes
-        :param kwargs: Additional User Pool attributes
         :return response: Response from Cognito
 
         Example response::
@@ -299,16 +297,18 @@ def register(self, username, password, attr_map=None):
             }
         }
         """
-
-        attributes= dict(self.base_attributes.items() + self.custom_attributes.items())
-        cognito_attributes = dict_to_cognito(attributes,attr_map)
-        user_attrs = [{'Name': key, 'Value': value} for key, value in attributes.items()]
-        response = self.client.sign_up(
-            ClientId=self.client_id,
-            Username=username,
-            Password=password,
-            UserAttributes= cognito_attributes
-        )
+        attributes = self.base_attributes.copy()
+        if self.custom_attributes:
+            attributes.update(self.custom_attributes)
+        cognito_attributes = dict_to_cognito(attributes, attr_map)
+        params = {
+            'ClientId': self.client_id,
+            'Username': username,
+            'Password': password,
+            'UserAttributes': cognito_attributes
+        }
+        self._add_secret_hash(params, 'SecretHash')
+        response = self.client.sign_up(**params)
 
         attributes.update(username=username, password=password)
         self._set_attributes(response, attributes)
@@ -340,11 +340,11 @@ def confirm_sign_up(self,confirmation_code,username=None):
         """
         if not username:
             username = self.username
-        self.client.confirm_sign_up(
-            ClientId=self.client_id,
-            Username=username,
-            ConfirmationCode=confirmation_code
-        )
+        params = {'ClientId': self.client_id,
+                  'Username': username,
+                  'ConfirmationCode': confirmation_code}
+        self._add_secret_hash(params, 'SecretHash')
+        self.client.confirm_sign_up(**params)
 
     def admin_authenticate(self, password):
         """
@@ -356,11 +356,7 @@ def admin_authenticate(self, password):
                 'USERNAME': self.username,
                 'PASSWORD': password
             }
-        if self.client_secret is not None:
-            auth_params.update({
-                'SECRET_HASH':
-                AWSSRP.get_secret_hash(self.username, self.client_id,
-                                       self.client_secret)})
+        self._add_secret_hash(auth_params, 'SECRET_HASH')
         tokens = self.client.admin_initiate_auth(
             UserPoolId=self.user_pool_id,
             ClientId=self.client_id,
@@ -552,9 +548,7 @@ def renew_access_token(self):
         Sets a new access token on the User using the refresh token.
         """
         auth_params = {'REFRESH_TOKEN': self.refresh_token}
-        if self.client_secret is not None:
-            auth_params.update({'SECRET_HASH':
-                                AWSSRP.get_secret_hash(self.username, self.client_id, self.client_secret)})
+        self._add_secret_hash(auth_params, 'SECRET_HASH')
         refresh_response = self.client.initiate_auth(
             ClientId=self.client_id,
             AuthFlow='REFRESH_TOKEN',
@@ -574,10 +568,12 @@ def initiate_forgot_password(self):
         """
         Sends a verification code to the user to use to change their password.
         """
-        self.client.forgot_password(
-            ClientId=self.client_id,
-            Username=self.username
-        )
+        params = {
+            'ClientId': self.client_id,
+            'Username': self.username
+        }
+        self._add_secret_hash(params, 'SecretHash')
+        self.client.forgot_password(**params)
 
 
     def delete_user(self):
@@ -601,12 +597,13 @@ def confirm_forgot_password(self, confirmation_code, password):
         to retrieve a forgotten password
         :param password: New password
         """
-        response = self.client.confirm_forgot_password(
-            ClientId=self.client_id,
-            Username=self.username,
-            ConfirmationCode=confirmation_code,
-            Password=password
-        )
+        params = {'ClientId': self.client_id,
+                  'Username': self.username,
+                  'ConfirmationCode': confirmation_code,
+                  'Password': password
+                  }
+        self._add_secret_hash(params, 'SecretHash')
+        response = self.client.confirm_forgot_password(**params)
         self._set_attributes(response, {'password': password})
 
     def change_password(self, previous_password, proposed_password):
@@ -621,6 +618,16 @@ def change_password(self, previous_password, proposed_password):
         )
         self._set_attributes(response, {'password': proposed_password})
 
+    def _add_secret_hash(self, parameters, key):
+        """
+        Helper function that computes SecretHash and adds it
+        to a parameters dictionary at a specified key
+        """
+        if self.client_secret is not None:
+            secret_hash = AWSSRP.get_secret_hash(self.username, self.client_id,
+                                                 self.client_secret)
+            parameters[key] = secret_hash
+
     def _set_attributes(self, response, attribute_dict):
         """
         Set user attributes based on response code

warrant/tests/tests.py

@@ -90,15 +90,16 @@ def test_verify_token(self):
     #     self.assertEqual(self.user.access_token,None)
 
     @patch('warrant.Cognito', autospec=True)
-    def test_register(self,cognito_user):
+    def test_register(self, cognito_user):
         u = cognito_user(self.cognito_user_pool_id, self.app_id,
                          username=self.username)
-        res = u.register('sampleuser','sample4#Password',
-                given_name='Brian',family_name='Jones',
-                name='Brian Jones',
-                email='[email protected]',
-                phone_number='+19194894555',gender='Male',
-                preferred_username='billyocean')
+        u.add_base_attributes(
+            given_name='Brian', family_name='Jones',
+            name='Brian Jones', email='[email protected]',
+            phone_number='+19194894555', gender='Male',
+            preferred_username='billyocean')
+        res = u.register('sampleuser', 'sample4#Password')
+
         #TODO: Write assumptions