browser-use-oauth/lib/agents/get_sso_list.py
tv0924@icloud.com 495b3a52da [Update]
2025-06-22 22:19:30 +09:00

66 lines
2.4 KiB
Python

import json
from pydantic import BaseModel
from browser_use import (
Agent,
Controller,
)
from lib.agents.run_agent import run_agent
from lib.utils.logger import logger
from lib.browser_use_utils.create_google_ai import create_google_ai
from lib.config import GOOGLE_MODEL, GOOGLE_PLANNER_MODEL
NOT_FOUND_SSO_LIST = 0
FOUND_SSO_LIST = 1
class EachSSOProvider(BaseModel):
provider: str
oauth_uri: str | None = None
class FindLoginPageResponse(BaseModel):
EachSSOProviders: list[EachSSOProvider] | None = None
status: int = NOT_FOUND_SSO_LIST # 0 if not found,
msg: str | None = None
async def get_sso_list(target_url, session) -> tuple[bool, str | None]:
initial_actions = [{"open_tab": {"url": target_url}}]
task = "Navigate to the login page, and return the result in the specified format."
extend_planner_system_message = """
You are an expert in finding login pages.
Your task is to navigate to the login page of the given URL.
Once you reach the login page, stop and return a JSON object that matches the following schema:
```json
{
"status": 1, # 1 if login page found, 0 otherwise
"url": "https://example.com/login" # Full URL of the login page if found
}
Return only this JSON object. Do not include any explanation or additional text.
"""
controller = Controller(output_model=FindLoginPageResponse, exclude_actions=['search_google'])
agent = Agent(
browser_session=session,
initial_actions=initial_actions,
task=task,
llm=create_google_ai(GOOGLE_MODEL),
planner_llm=create_google_ai(GOOGLE_PLANNER_MODEL),
controller=controller,
extend_planner_system_message=extend_planner_system_message,
)
is_failed, final_result = await run_agent(agent)
if is_failed:
logger(f"⚠️ 스캔 실패: {target_url} | {final_result}")
print(f"⚠️ 스캔 실패: {target_url} | {final_result}")
return False, None;
data = json.loads(final_result)
try:
resp = FindLoginPageResponse(**data)
if resp.status == FOUND_SSO_LIST:
return True, resp
else:
return False, None
except Exception as e:
logger(f"⚠️ 결과 파싱 실패: {target_url} | {e}\n원본 결과: {data.msg}")
print(f"⚠️ 결과 파싱 실패: {target_url} | {e}\n원본 결과: {data.msg}")
return False, data.msg