mirror of
https://github.com/browser-use/browser-use
synced 2026-05-06 17:52:15 +02:00
6a5c0b61f2
This commit introduces a new script, `2fa.py`, which implements a custom action to retrieve 2FA/MFA codes using the `pyotp` library. The action extracts the OTP secret key from an environment variable and generates a time-based one-time password. Additionally, the script includes an example task demonstrating how to use the 2FA action within an agent-based workflow. ### Key Changes: - New file `examples/custom-functions/2fa.py` added. - Custom action `get_otp_2fa` to generate 2FA codes. - Example usage of the action in an asynchronous context. This addition enhances the functionality of the agent system by providing secure 2FA code retrieval capabilities.
68 lines
1.9 KiB
Python
68 lines
1.9 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
import sys
|
|
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
import pyotp
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
from browser_use import ActionResult, Agent, Controller
|
|
|
|
# Set up logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
controller = Controller()
|
|
|
|
|
|
@controller.registry.action('Get 2FA code from when OTP is required')
|
|
async def get_otp_2fa() -> ActionResult:
|
|
"""
|
|
Custom action to retrieve 2FA/MFA code from OTP secret key using pyotp.
|
|
The OTP secret key should be set in the environment variable OTP_SECRET_KEY.
|
|
"""
|
|
secret_key = os.environ.get('OTP_SECRET_KEY')
|
|
if not secret_key:
|
|
raise ValueError('OTP_SECRET_KEY environment variable is not set')
|
|
|
|
totp = pyotp.TOTP(secret_key, digits=6)
|
|
code = totp.now()
|
|
return ActionResult(extracted_content=code)
|
|
|
|
|
|
async def main():
|
|
# Example task using the 1Password 2FA action
|
|
task = """
|
|
Steps:
|
|
1. Go to https://authenticationtest.com/totpChallenge/ and try to log in.
|
|
2. If prompted for 2FA code:
|
|
2.1. Use the get_2fa_code action to retrieve the 2FA code.
|
|
2.2. Submit the code provided by the get_2fa_code action.
|
|
|
|
Considerations:
|
|
- ALWAYS use the get_2fa_code action to retrieve the 2FA code if needed.
|
|
- NEVER skip the 2FA step if the page requires it.
|
|
- NEVER extract the code from the page.
|
|
- NEVER use a code that is not generated by the get_2fa_code action.
|
|
- NEVER hallucinate the 2FA code, always use the get_2fa_code action to get it.
|
|
|
|
You are completely FORBIDDEN to use any other method to get the 2FA code.
|
|
"""
|
|
|
|
model = ChatOpenAI(model='gpt-4o')
|
|
agent = Agent(task=task, llm=model, controller=controller)
|
|
|
|
result = await agent.run()
|
|
print(f'Task completed with result: {result}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(main())
|