In my previous post, I discussed how to use Streamlit for building forms and enabling users to upload images. This project also needed an additional layer of access control to only allow a subset of users with referral codes to view the form.
The main issue to address here is to prevent brute-force attacks on the referral code input field. However, this could extend to any resource in Streamlit that you may want to rate-limit to not overwhelm the backend services.
Streamlit has a cool feature called 'Session State' which allows us to keep track of user actions. According to Streamlit's documentation:
"Session State is a way to share variables between reruns, for each user session. In addition to the ability to store and persist state, Streamlit also exposes the ability to manipulate state using Callbacks. Session state also persists across apps inside a multipage app."
Using this feature we could keep track of how many times a user attempts to enter the referral code and provide immediate feedback on incorrect attempts so users can use the appropriate channel to get a valid referral code.
If the user continues to brute force the referral code, we will employ a Lambda trigger to disable the user account in the AWS Cognito user pool.
Example Streamlit code which explains how to use 'session state' to track user re-tries.
import streamlit as st
import json
if 'count' not in st.session_state:
st.session_state.count = 0
def increment_counter():
st.session_state.count += 1
st.write("**Please enter the invite code.**")
if st.session_state.count <= 5:
ref_code = st.text_input("Enter the referral code", on_change=increment_counter, key="ref_code", help="You have 5 retries to enter the correct referral code.")
else:
st.error("You have exceeded the maximum number of retries.")
ref_code = None
disable_user_api = "https://test.execute-api.us-east-2.amazonaws.com/disable_user"
disable_user_payload = {
"uuid": uuid
}
disable_user_headers = {
'Content-Type': 'application/json',
'Authorization': access_token
}
try:
disable_user_response = requests.post(disable_user_api, data=json.dumps(disable_user_payload), headers=disable_user_headers)
st.success("Your account has been disabled. Please contact the administrator.")
except Exception as disable_user_err:
st.error(f'An error occurred: {disable_user_err}')
if ref_code == "1234":
st.write("Referral code is correct. Please enter your details below.")
else:
st.error("Referral code is invalid. If you do not have a referral code, please write to us @ some email.")
st.error("Invalid attempt count: " + str(st.session_state.count) + ". You have " + str(5 - st.session_state.count) + " attempts left.")
Streamlit renders the above text box upon running the Python code streamlit run app.py. After submitting an incorrect code, Streamlit throws the following error message:
In this example, I've set the maximum retries to 5 before the account gets disabled. Session State uses the callback functionality by using the 'on_change' parameter to keep track of the number of submission attempts.
The 6th incorrect submission will trigger the backend Lambda function by sending an HTTP POST request along with the user's unique account ID (uuid) in the request body.
Example Lambda function:
import json
import boto3
def lambda_handler(event, context):
http_method = event['requestContext']['http']['method']
if http_method == 'POST':
body = json.loads(event['body'])
uuid = body['uuid']
disable_user(uuid)
handle_sns(uuid)
return {
'statusCode': 200,
'body': json.dumps("User disabled.")
}
if http_method == 'GET':
return {
'statusCode': 200
}
return {
'statusCode': 200
}
def disable_user(uuid):
client_idp = boto3.client('cognito-idp')
client_idp.admin_disable_user(
UserPoolId='us-east-1_xxxxxxx',
Username=uuid
)
def handle_sns(uuid):
client_sns = boto3.client('sns')
sns_message = """
User account has been disabled. Find the UUID below:
UUID : {uuid}
""".format(uuid=uuid)
client_sns.publish(
TopicArn='arn:aws:sns:us-east-1:XXXXXXXXXX:topic-name',
Message= sns_message,
Subject='Account disabled')
The example lambda function performs two actions:
Uses boto3.client('cognito-idp') to disable the user account in AWS Cognito based on the User-Pool ID.
Uses boto3.client('sns') to send a notification to the administrator's email as per the TopicARN (SNS configuration).
For the lambda function to properly execute the actions, ensure that the AWS role associated with the lambda function has permission to disable the user for the particular AWS pool ID as an administrator and to trigger the SNS Topic.
