Legacy Lobotomy — Fixing Vulnerabilities in the Profile Management API

The photo was generated by the author using ChatGPT-4.

This is the 8th tutorial in the series about refactoring a legacy Django project, where we are going to talk about fixing vulnerabilities found in the profile management API described in the previous tutorial. The main goal of the fixing process is to ensure everything that used to work still works and that we solve the problems we found. So, after we’re done, the system should work better than it did before.

To get more ideas regarding the project used for this tutorial or check other tutorials in this series, check the introductory article published earlier.

Legacy Lobotomy — Confident Refactoring of a Django Project

Origin branch and destination branch

If you want to follow the steps described in this tutorial, you should start from the profile-management-api-analysis branch.If you want to only see final code state after all changes from this tutorial are applied, you can check the profile-management-api-fixing-vulnerabilities branch.

Preparation

To start working on tests for the profile management API, we need to make some preparations.

Extend the UserFactory class with extra fields

In one of the previous tutorials we have added UserFactory class which allows us to generate users for using in tests. But this class allows us to create only the simplest user, with email and password fields filled. The rest fields are either not set or have default values. In order to test profile management endpoints properly we should add some randomness to our data. Let’s extend the UserFactory class with more fields:

class UserFactory(factory.django.DjangoModelFactory):
# ———- Other properties are omitted for simplicity ———— #
first_name = factory.Faker(‘first_name’)
last_name = factory.Faker(‘last_name’)
age = factory.Faker(‘pyint’, min_value=13, max_value=99)
gender = factory.Faker(‘random_element’, elements=[item[0] for item in User.GENDER_CHOICES])
guardian_email = factory.Faker(’email’)
accepted_terms_cond = factory.Faker(‘pybool’)
activity = factory.Faker(‘random_element’, elements=[item[0] for item in User.ACTIVITY_CHOICES])
total_points = factory.Faker(‘pyint’)
first_login = factory.Faker(‘pybool’)
is_superuser = False # by default this factory should create regular users

If we run tests, they all should pass.

Prepare a class for testing profile management API

In the previous tutorial we have added test_users_viewset.py module with a set of tests to demonstrate vulnerabilities of the current implementation of the profile management API. Let’s apply some changes to this module:

Step 1. Rename it to test_profile_management.py

Step 2. Change the class name from TestUsersViewSet to TestProfileManagement

Step 3. Add two helper methods at the beginning of the TestProfileManagement class as shown below

class TestProfileManagement:
def user_details_url(self, user):
return f’/api/users/{user.pk}/’

def update_user_payload(self, user):
return {
’email’: user.email,
‘first_name’: user.first_name,
‘last_name’: user.last_name,
‘age’: user.age,
‘gender’: user.gender,
‘guardian_email’: user.guardian_email,
‘accepted_terms_cond’: user.accepted_terms_cond,
‘activity’: user.activity,
‘total_points’: user.total_points,
‘first_login’: user.first_login,
‘is_superuser’: user.is_superuser,
}

We have added user_details_url method which returns a URL for accessing the currently signed-in user details. This method will be removed later, when we get rid of the UserViewSet class and replace it with another view class.

Also, there is update_user_payload method which returns a dictionary object with all fields defined in the UserSerializer class except for pk. We will be using this helper method to build a default payload for the PUT method in tests.

Writing tests and fixing issues with current implementation

In this section, I will be implementing several test groups to evaluate the expected behavior. Each time a test fails, I will modify the corresponding code to ensure the successful completion of the tests.

Anonymous user access tests

As the tested API is designed for retrieving and modifying user details, it is evident that none of the endpoints should be accessible to anonymous users. Let’s include the following parametrized test to verify if the API functions as expected.

class TestProfileManagement:
# ———- Other properties and methods are omitted for simplicity ———— #
@pytest.mark.django_db
@pytest.mark.parametrize(‘method’, [‘get’, ‘put’, ‘delete’])
def test_when_user_is_not_authenticated_then_401_unauthorized_code_should_be_returned(self, method, user, client):
method_callable = getattr(client, method)
response = method_callable(self.user_details_url(user))

assert response.status_code == status.HTTP_401_UNAUTHORIZED

In this test we want to make HTTP client to send GET, PUT and DELETE requests to the user details endpoint. We expect that all three calls will result in a 401 Unauthorized HTTP code. We don’t pass any payload with the PUT method. This omission is intentional, as the server’s validation process checks if the request is made on behalf of an authorized user before validating the request payload.

If we run tests, they all should pass. It means that all three HTTP methods are accessible and protected from the anonymous access.

Authenticated user access tests

Earlier we discussed that the profile management API should enable the retrieval of user details for both ordinary users and superusers. However, only ordinary users should have the capability to modify their details and delete their profiles. Let’s add the following tests to cover the described behavior.

class TestProfileManagement:
# ———- Other properties and methods are omitted for simplicity ———— #
@pytest.mark.django_db
@pytest.mark.parametrize(‘user__is_superuser’, [
pytest.param(True, id=’user is a superuser’),
pytest.param(False, id=’user is an ordinary user’),
])
def test_get_method_with_authenticated_user(self, user, auth_client):
response = auth_client.get(self.user_details_url(user))

assert response.status_code == status.HTTP_200_OK
assert response.json() == {
‘pk’: user.pk,
’email’: user.email,
‘first_name’: user.first_name,
‘last_name’: user.last_name,
‘age’: user.age,
‘gender’: user.gender,
‘guardian_email’: user.guardian_email,
‘accepted_terms_cond’: user.accepted_terms_cond,
‘activity’: user.activity,
‘total_points’: user.total_points,
‘first_login’: user.first_login,
‘is_superuser’: user.is_superuser,
}

In this test we verify that ordinary users and superusers can request the details of their profiles and can see correct information. Here we specified user__is_superuser parameter which should change for each test run. This way we can override is_superuser property of the user instance created by the user fixture. This cool feature is provided by the pytest-factoryboy package.

class TestProfileManagement:
# ———- Other properties and methods are omitted for simplicity ———— #
@pytest.mark.django_db
@pytest.mark.parametrize(‘user__is_superuser, expected_status’, [
pytest.param(True, status.HTTP_403_FORBIDDEN, id=’a superuser should not be able to update the profile’),
pytest.param(False, status.HTTP_200_OK, id=’an ordinary user should be able to update the profile’),
])
def test_put_method_with_authenticated_user(self, expected_status, user, auth_client):
response = auth_client.put(self.user_details_url(user), self.update_user_payload(user))

assert response.status_code == expected_status

With this test we verify that only ordinary users can update their profiles. Here we use two parameters, user__is_superuser and expected_status. The last one contains the expected status which the endpoint should return for each specific test case. With using the update_user_payload method we build a default payload for the PUT method.

class TestProfileManagement:
# ———- Other properties and methods are omitted for simplicity ———— #
@pytest.mark.django_db
@pytest.mark.parametrize(‘user__is_superuser, expected_status’, [
pytest.param(True, status.HTTP_403_FORBIDDEN, id=’a superuser should not be able to delete the profile’),
pytest.param(False, status.HTTP_204_NO_CONTENT, id=’an ordinary user should be able to delete the profile’),
])
def test_delete_method_with_authenticated_user(self, expected_status, user, auth_client):
response = auth_client.delete(self.user_details_url(user))

assert response.status_code == expected_status

This test is very similar to the previous one, but intended for verifying that only ordinary users can delete their profiles.

If we run the tests now again, we will see that the subtest user is a superuser of the test_get_method_with_authenticated_user test doesn’t pass. The reason for this is the permissions defined for the whole UserViewSet class, so they are the same for all endpoints provided by the viewset class.

Fixing problems with retrieving profile details for superusers

To make this test pass, we need to check different permissions depending on the action triggered by the request. In order to do this, we need:

Step 1. Remove the following property from the UserViewSet class:

permission_classes = (IsAuthenticated, IsNotSuperuser)

Step 2. Add the following method to the UserViewSet class:

class UserViewSet(viewsets.ModelViewSet):
# ———- Other properties and methods are omitted for simplicity ———— #
def get_permissions(self):
if self.action == ‘retrieve’:
self.permission_classes = (IsAuthenticated,)
else:
self.permission_classes = (IsAuthenticated, IsNotSuperuser)

return super().get_permissions()

This method is defined in the APIView base class and its code looks as follows.

class APIView:
# ———- Other properties and methods are omitted for simplicity ———— #
def get_permissions(self):
return [permission() for permission in self.permission_classes]

As we can see, the original method just loops over the permission_classes collection and creates an instance of each permission class. Our implementation assigns (IsAuthenticated,) value to this property if the retrieve action is triggered. Otherwise, (IsAuthenticated, IsNotSuperuser) value is assigned. Once the permission_classes property is set up, we call the method of the parent class.

Now all tests should pass.

Fixing the issue of accessing other users’ details

In order to prevent access to the details of other users, we can restrict a queryset used by UserViewSet class to current user only. Let’s add a get_queryset method as shown below:

class UserViewSet(viewsets.ModelViewSet):
# ———- Other properties and methods are omitted for simplicity ———— #
def get_queryset(self):
return User.objects.filter(pk=self.request.user.id)

If we run tests, we will see that 4 tests of the TestProfileManagement class fail. Our goal is to reverse the failed tests to ensure that the identified issues no longer exist.

The test test_any_authenticated_user_can_see_the_list_of_all_users_in_the_system fails because we expect the GET /api/users/ endpoint to return a list of details for all users in the system (the current user and another user in this case). However, it only returned a list of details for the current user. Let’s modify this test in the following way:

class TestProfileManagement:
# ———- Other methods are omitted for simplicity ———— #
@pytest.mark.django_db
def test_get_users_method_should_return_only_current_user_details(self, auth_client, user, user_factory):
another_user = user_factory()

response = auth_client.get(‘/api/users/’)
returned_users = {item[‘pk’] for item in response.json()}

assert response.status_code == status.HTTP_200_OK
assert returned_users == {user.pk}

Here is the list of changes made to the test:

The test was renamed to test_get_users_method_should_return_only_current_user_details to properly reflect its new meaning — we expect this endpoint to return only the list of details of the current userThe line assert returned_users == {user.pk, another_user.pk} was replaced with assert returned_users == {user.pk} to check that only current user’s details are returned by the GET /api/users/ method

The test test_any_authenticated_user_can_see_details_of_any_user_in_the_system fails because the GET /api/users/{user_id}/ endpoint returns data only if the value of the user_id parameter is equal to self.request.user.id. Otherwise, it returns a 404 Not Found HTTP error. In this test, we checked that any authenticated user can see details of any other user in the system. After we restricted the queryset used by the UserViewSet class, a user cannot request details of other users anymore. Now, a user can request only their own details. Let’s modify this test in the following way:

class TestProfileManagement:
# ———- Other methods are omitted for simplicity ———— #
@pytest.mark.django_db
def test_authenticated_user_cannot_see_details_of_any_other_user_in_the_system(self, auth_client, user_factory):
another_user = user_factory()

response = auth_client.get(f’/api/users/{another_user.pk}/’)

assert response.status_code == status.HTTP_404_NOT_FOUND

Here is the list of changes made to the test:

The test was renamed to test_authenticated_user_cannot_see_details_of_any_other_user_in_the_system to emphasize that obtaining details of other users is prohibitedThe line assert response.status_code == status.HTTP_200_OK was replaced with assert response.status_code == status.HTTP_404_NOT_FOUND to ensure the server returns a proper HTTP codeThe line assert response.json()[‘pk’] == another_user.pk was removed because now it doesn’t make sense

The test test_any_authenticated_user_can_make_another_user_a_superuser also fails because an authenticated user attempts to change details of another user. Let’s change the test as shown below:

class TestProfileManagement:
# ———- Other methods are omitted for simplicity ———— #
@pytest.mark.django_db
def test_authenticated_user_cannot_make_another_user_a_superuser(self, auth_client, user_factory):
# first_login=False is used to work around a bug with attempt to change request data
another_user = user_factory(is_superuser=False, first_login=False)
response = auth_client.patch(f’/api/users/{another_user.pk}/’, {‘is_superuser’: True})

another_user.refresh_from_db()

assert response.status_code == status.HTTP_404_NOT_FOUND
assert another_user.is_superuser is False

Here is the list of changes made to the test:

The test was renamed to test_authenticated_user_cannot_make_another_user_a_superuser to show that we expect that one user cannot make another user a superuser anymoreThe line assert response.status_code == status.HTTP_200_OK was replaced with assert response.status_code == status.HTTP_404_NOT_FOUND to ensure the server returns a proper HTTP codeWe replaced lines that verified that another_user is a superuser with the line assert another_user.is_superuser is False that verifies that the value of the is_superuser flag of the another_user wasn’t changed with the request

The test test_any_authenticated_user_can_delete_another_user fails due to the same reason — now one user cannot delete another user’s profile. Let’s update the test to reflect expected behavior:

class TestProfileManagement:
# ———- Other methods are omitted for simplicity ———— #
@pytest.mark.django_db
def test_authenticated_user_cannot_delete_another_user(self, auth_client, user_factory):
another_user = user_factory()

response = auth_client.delete(f’/api/users/{another_user.pk}/’)

assert response.status_code == status.HTTP_404_NOT_FOUND
assert User.objects.filter(pk=another_user.pk).exists() is True

Here is the list of changes made to the test:

The test was renamed to test_authenticated_user_cannot_delete_another_user to show that we expect that one user cannot delete another userThe line assert response.status_code == status.HTTP_200_OK was replaced with assert response.status_code == status.HTTP_404_NOT_FOUND to ensure the server returns a proper HTTP codeInstead of checking that another_user doesn’t exist anymore, we check that the record still exists by adding the line assert User.objects.filter(pk=another_user.pk).exists() is True

If we run the tests now, all they should pass. Everything seems good now. We know that a user cannot see and update details and delete profiles of other users. But there is still an issue that a user can make themselves a superuser that proves the test test_any_authenticated_user_can_make_themselves_a_superuser. Let’s keep this test as it is for now because the issue it highlights is deeper, and we need to introduce more changes to address it.

Verifying the first_login property is updated correctly

If we look at the update method of the UserViewSet class we will notice that there is custom logic for updating the first_login property of the user. This flag was added to the system to force a user to update their profile information after their very first sign-in to the system.

class UserViewSet(viewsets.ModelViewSet):
# ———- Other methods are omitted for simplicity ———— #
def update(self, request, *args, **kwargs):
data = request.data
instance = self.get_object()
if instance.first_login:
data[‘first_login’] = False
serializer = UserSerializer(instance, data=data, partial=True)
# ———- The rest of the method is omitted for simplicity ———— #

Let’s add a test for checking different situations with changing this field:

If the first_login property of the updated user is True, it should become FalseIf the first_login property of the updated user is False, it should remain unchangedThe result shouldn’t depend on the value passed for the first_login fieldThe result shouldn’t depend on whether the first_login field is passed with the request or notclass TestProfileManagement:
# ———- Other methods are omitted for simplicity ———— #
@pytest.mark.django_db
@pytest.mark.parametrize(‘updated_first_login’, [
pytest.param(True, id=’updated_first_login = True’),
pytest.param(False, id=’updated_first_login = False’),
])
@pytest.mark.parametrize(‘user__first_login’, [
pytest.param(True, id=’user__first_login = True’),
pytest.param(False, id=’user__first_login = False’),
])
def test_put_method_first_login_value_should_become_false_regardless_of_the_request_value(
self, user, auth_client, updated_first_login
):
payload = {**self.update_user_payload(user), ‘first_login’: updated_first_login}
auth_client.put(self.user_details_url(user), payload)

user.refresh_from_db()

assert user.first_login is False

This test covers requirements 1–3. It creates a user with first_login = True and first_login = False, and calls the PUT method with passing values True and False for the first_login field.

We also need to add another test that checks the case when the first_login property of the updated user is either True or False and such a field isn’t passed with the request.

class TestProfileManagement:
# ———- Other methods are omitted for simplicity ———— #
@pytest.mark.django_db
@pytest.mark.parametrize(‘user__first_login’, [
pytest.param(True, id=’user__first_login = True’),
pytest.param(False, id=’user__first_login = False’),
])
def test_put_method_first_login_value_should_become_false_even_if_request_does_not_contain_first_login_field(
self, user, auth_client
):
payload = self.update_user_payload(user)
payload.pop(‘first_login’)
auth_client.put(self.user_details_url(user), payload)

user.refresh_from_db()

assert user.first_login is False

If we run all the tests now, we will see that the subtest user__first_login = False-updated_first_login = True of the test test_put_method_first_login_value_should_become_false_regardless_of_the_request_value doesn’t pass because the test expects that after the update a user’s first_login field will be False, but it was True as it was passed with the request. So this test shows that it’s possible to override this field via API, but it shouldn’t be.

As we now have tests to cover all possible combinations of updating the first_login field value, we can change the code to make the tests pass. Since we want this field to be False regardless of the current and passed values, we can replace this code

class UserViewSet(viewsets.ModelViewSet):
# ———- Other methods are omitted for simplicity ———— #
def update(self, request, *args, **kwargs):
data = request.data
instance = self.get_object()
if instance.first_login:
data[‘first_login’] = False
serializer = UserSerializer(instance, data=data, partial=True)
# ———- The rest of the method is omitted for simplicity ———— #

with the code below

class UserViewSet(viewsets.ModelViewSet):
# ———- Other methods are omitted for simplicity ———— #
def update(self, request, *args, **kwargs):
data = {**request.data, ‘first_login’: False}
serializer = UserSerializer(self.get_object(), data=data, partial=True)
# ———- The rest of the method is omitted for simplicity ———— #

Now, when we run the tests, all of them should pass. Instead of modifying the original data object, we create a copy of it and set the value of the first_login field to False. We don’t need to check the actual value of this field because it doesn’t make sense to do so, as it should be set to False after the first update and then never change.

Fixing the issue of a user being able to make themselves a superuser

Although all tests pass and some issues have already been fixed, users are still able to change some properties that are supposed to be unchangeable. In this section we are going to make sure the users can update their profiles and make some fields readonly.

We’ll focus on supporting only the PUT method since the app doesn’t use PATCH. When we examine the instantiation of the UserSerializer class in the update method, we notice that both PUT and PATCH requests supported and function identically. This is because partial=True instruction indicates a partial update, allowing some fields to be absent. For a PUT request, we expect all fields to be included in the request, and if any are missing, a validation error should be generated.

serializer = UserSerializer(self.get_object(), data=data, partial=True)

Also, we need to understand which fields users should be able to update. Based on the app’s code, these fields are:

first_namelast_nameagegenderguardian_emailactivity

The following fields should be readonly:

pk – this is a primary key field, it’s readonly by defaultemail – creating new users and changing the email address is delegated to admins in the system, it’s a specific requirementaccepted_terms_cond – there is another API endpoint for accepting terms and conditionstotal_points – a user cannot change the total number of earned points — this is a quite serious issuefirst_login – this field is managed on the server, users shouldn’t be able to change the value of this fieldis_superuser – obviously, a user shouldn’t be able to make themselves a superuser

Let’s add the following test to check that a user cannot update these fields via the API. Since we manage the value of the first_login field ourselves and don’t allow it to be overridden via an API call, and we have already implemented a set of tests for this field change, we need to check how the system behaves with other fields.

class TestProfileManagement:
# ———- Other methods are omitted for simplicity ———— #
@pytest.mark.django_db
@pytest.mark.parametrize(‘field_name, current_value, new_value’, [
pytest.param(‘id’, 100, 200, id=’id: 100 -> 200′),
pytest.param(’email’, ‘fake-user.old-email@fake.com’, ‘fake-user.new-email@fake.com’, id=’email’),
pytest.param(‘accepted_terms_cond’, False, True, id=’accepted_terms_cond: False -> True’),
pytest.param(‘accepted_terms_cond’, True, False, id=’accepted_terms_cond: True -> False’),
pytest.param(‘total_points’, 0, 100, id=’total_points’),
pytest.param(‘is_superuser’, False, True, id=’is_superuser: False -> True’),
pytest.param(‘is_superuser’, True, False, id=’is_superuser: True -> False’),
])
def test_put_method_readonly_fields_should_not_be_updated(
self, field_name, current_value, new_value, user_factory, auth_client_factory
):
user = user_factory(**{field_name: current_value})
payload = {**self.update_user_payload(user), field_name: new_value}

auth_client = auth_client_factory(user)
auth_client.put(self.user_details_url(user), payload)

user.refresh_from_db()

assert getattr(user, field_name) == current_value

This test is parametrized with field_name, current_value and new_value parameters. Every time we create a user with the current_value value of the property field_name, try to update this field with the new_value value via API method and verify that the value wasn’t changed. If we run tests now, we will see that 5 tests fail. It means that not all of these readonly fields are actually readonly. The test is_superuser: True -> False passes because the server returns 403 Forbidden HTTP error since superusers are not allowed to change their profile information (check IsNotSuperuser permission class for more details). The test id: 100 -> 200 passes because the id field is a primary key and readonly by default.

To make other tests pass and fix the problem with updating readonly fields, we need to update the UserSerializer class with the list of fields which should be readonly. There are several options to do this, but the simplest one for our case is to add read_only_fields property into UserSerializer.Meta class:

class UserSerializer(serializers.ModelSerializer):
class Meta:
# ———- Other properties are omitted for simplicity ———— #
read_only_fields = (
’email’,
‘accepted_terms_cond’,
‘total_points’,
‘first_login’,
‘is_superuser’
)

If we run tests now, we will see that 4 of them fail.

In the selected tests a user was created with the True value of the first_login field, and we expected it to be False after API method call, but it remained True after we defined the list of readonly fields. The problem is that we added first_login field into this list as well, but we update its value by passing data = {**request.data.dict(), ‘first_login’: False} object into the serializer. Since this field is defined as readonly, serializer ignores it. To fix this problem we need to change this field another way. Fortunately, we have a simple way to do this.

Let’s remove data = {**request.data.dict(), ‘first_login’: False} declaration, replace data=data with data=request.data in the UserSerializer constructor and call the serializer’s save method with passing in the first_login parameter with False as a value. When we pass parameters directly into the serializer’s save method, we avoid the validation phase and such values will be used as is.

The final look of the update method of the UserViewSet class should look as follows:

class UserViewSet(viewsets.ModelViewSet):
# ———- Other properties and methods are omitted for simplicity ———— #
def update(self, request, *args, **kwargs):
serializer = UserSerializer(self.get_object(), data=request.data, partial=True)
serializer.is_valid(raise_exception=True)
serializer.save(first_login=False)
return Response(serializer.data)

After this update, all tests but test_any_authenticated_user_can_make_themselves_a_superuser should pass. If we investigate details, we will see that it failed because it expected the is_superuser flag to be changed from False to True, but it wasn’t. It means that a user cannot make them a superuser anymore.

We can safely remove the test_any_authenticated_user_can_make_themselves_a_superuser test because we already have a test that checks impossibility to change is_superuser field, so it’s not needed anymore.

Also, we can safely remove all readonly fields from the update_user_payload method and make it simpler:

class TestProfileManagement:
# ———- Other methods are omitted for simplicity ———— #
def update_user_payload(self, user):
return {
‘first_name’: user.first_name,
‘last_name’: user.last_name,
‘age’: user.age,
‘gender’: user.gender,
‘guardian_email’: user.guardian_email,
‘activity’: user.activity,
}

After that we need to remove payload.pop(‘first_login’) line from the test_put_method_first_login_value_should_become_false_even_if_request_does_not_contain_first_login_field test method.

Let’s run all tests again and make sure they all pass.

Conclusion

Fixing vulnerabilities in the Profile Management API is crucial for ensuring the security and efficiency of our system. By carefully identifying and fixing these issues, we can protect user data and improve the overall reliability of the system. However, without a good set of automated tests, we risk introducing new bugs and spending more time on manual testing. In the next tutorial, we will discuss problems with validation in the Profile Management API and continue working on improving the quality of the project.

Legacy Lobotomy — Fixing Vulnerabilities in the Profile Management API was originally published in Level Up Coding on Medium, where people are continuing the conversation by highlighting and responding to this story.

​ Level Up Coding – Medium

about Infinite Loop Digital

We support businesses by identifying requirements and helping clients integrate AI seamlessly into their operations.

Gartner
Gartner Digital Workplace Summit Generative Al

GenAI sessions:

  • 4 Use Cases for Generative AI and ChatGPT in the Digital Workplace
  • How the Power of Generative AI Will Transform Knowledge Management
  • The Perils and Promises of Microsoft 365 Copilot
  • How to Be the Generative AI Champion Your CIO and Organization Need
  • How to Shift Organizational Culture Today to Embrace Generative AI Tomorrow
  • Mitigate the Risks of Generative AI by Enhancing Your Information Governance
  • Cultivate Essential Skills for Collaborating With Artificial Intelligence
  • Ask the Expert: Microsoft 365 Copilot
  • Generative AI Across Digital Workplace Markets
10 – 11 June 2024

London, U.K.