Feat: Add box connector (#11845)

### What problem does this PR solve?

Feat: Add box connector

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Magicbook1108
2025-12-12 10:23:40 +08:00
committed by GitHub
parent a6bd765a02
commit 7db9045b74
19 changed files with 1019 additions and 131 deletions

View File

@ -28,7 +28,6 @@ from api.db.services import UserService
from api.utils.json_encode import CustomJSONEncoder
from api.utils import commands
from flask_mail import Mail
from quart_auth import Unauthorized
from common import settings
from api.utils.api_utils import server_error_response
@ -42,7 +41,6 @@ __all__ = ["app"]
app = Quart(__name__)
app = cors(app, allow_origin="*")
smtp_mail_server = Mail()
# Add this at the beginning of your file to configure Swagger UI
swagger_config = {

View File

@ -28,11 +28,12 @@ from api.db import InputType
from api.db.services.connector_service import ConnectorService, SyncLogsService
from api.utils.api_utils import get_data_error_result, get_json_result, get_request_json, validate_request
from common.constants import RetCode, TaskStatus
from common.data_source.config import GOOGLE_DRIVE_WEB_OAUTH_REDIRECT_URI, GMAIL_WEB_OAUTH_REDIRECT_URI, DocumentSource
from common.data_source.google_util.constant import GOOGLE_WEB_OAUTH_POPUP_TEMPLATE, GOOGLE_SCOPES
from common.data_source.config import GOOGLE_DRIVE_WEB_OAUTH_REDIRECT_URI, GMAIL_WEB_OAUTH_REDIRECT_URI, BOX_WEB_OAUTH_REDIRECT_URI, DocumentSource
from common.data_source.google_util.constant import WEB_OAUTH_POPUP_TEMPLATE, GOOGLE_SCOPES
from common.misc_utils import get_uuid
from rag.utils.redis_conn import REDIS_CONN
from api.apps import login_required, current_user
from box_sdk_gen import BoxOAuth, OAuthConfig, GetAuthorizeUrlOptions
@manager.route("/set", methods=["POST"]) # noqa: F821
@ -117,8 +118,6 @@ def rm_connector(connector_id):
return get_json_result(data=True)
GOOGLE_WEB_FLOW_STATE_PREFIX = "google_drive_web_flow_state"
GOOGLE_WEB_FLOW_RESULT_PREFIX = "google_drive_web_flow_result"
WEB_FLOW_TTL_SECS = 15 * 60
@ -129,10 +128,7 @@ def _web_state_cache_key(flow_id: str, source_type: str | None = None) -> str:
When source_type == "gmail", a different prefix is used so that
Drive/Gmail flows don't clash in Redis.
"""
if source_type == "gmail":
prefix = "gmail_web_flow_state"
else:
prefix = GOOGLE_WEB_FLOW_STATE_PREFIX
prefix = f"{source_type}_web_flow_state"
return f"{prefix}:{flow_id}"
@ -141,10 +137,7 @@ def _web_result_cache_key(flow_id: str, source_type: str | None = None) -> str:
Mirrors _web_state_cache_key logic for result storage.
"""
if source_type == "gmail":
prefix = "gmail_web_flow_result"
else:
prefix = GOOGLE_WEB_FLOW_RESULT_PREFIX
prefix = f"{source_type}_web_flow_result"
return f"{prefix}:{flow_id}"
@ -180,7 +173,7 @@ async def _render_web_oauth_popup(flow_id: str, success: bool, message: str, sou
}
)
# TODO(google-oauth): title/heading/message may need to reflect drive/gmail based on cached type
html = GOOGLE_WEB_OAUTH_POPUP_TEMPLATE.format(
html = WEB_OAUTH_POPUP_TEMPLATE.format(
title=f"Google {source.capitalize()} Authorization",
heading="Authorization complete" if success else "Authorization failed",
message=escaped_message,
@ -204,8 +197,8 @@ async def start_google_web_oauth():
redirect_uri = GMAIL_WEB_OAUTH_REDIRECT_URI
scopes = GOOGLE_SCOPES[DocumentSource.GMAIL]
else:
redirect_uri = GOOGLE_DRIVE_WEB_OAUTH_REDIRECT_URI if source == "google-drive" else GMAIL_WEB_OAUTH_REDIRECT_URI
scopes = GOOGLE_SCOPES[DocumentSource.GOOGLE_DRIVE if source == "google-drive" else DocumentSource.GMAIL]
redirect_uri = GOOGLE_DRIVE_WEB_OAUTH_REDIRECT_URI
scopes = GOOGLE_SCOPES[DocumentSource.GOOGLE_DRIVE]
if not redirect_uri:
return get_json_result(
@ -271,8 +264,6 @@ async def google_gmail_web_oauth_callback():
state_id = request.args.get("state")
error = request.args.get("error")
source = "gmail"
if source != 'gmail':
return await _render_web_oauth_popup("", False, "Invalid Google OAuth type.", source)
error_description = request.args.get("error_description") or error
@ -313,9 +304,6 @@ async def google_gmail_web_oauth_callback():
"credentials": creds_json,
}
REDIS_CONN.set_obj(_web_result_cache_key(state_id, source), result_payload, WEB_FLOW_TTL_SECS)
print("\n\n", _web_result_cache_key(state_id, source), "\n\n")
REDIS_CONN.delete(_web_state_cache_key(state_id, source))
return await _render_web_oauth_popup(state_id, True, "Authorization completed successfully.", source)
@ -326,8 +314,6 @@ async def google_drive_web_oauth_callback():
state_id = request.args.get("state")
error = request.args.get("error")
source = "google-drive"
if source not in ("google-drive", "gmail"):
return await _render_web_oauth_popup("", False, "Invalid Google OAuth type.", source)
error_description = request.args.get("error_description") or error
@ -391,3 +377,107 @@ async def poll_google_web_result():
REDIS_CONN.delete(_web_result_cache_key(flow_id, source))
return get_json_result(data={"credentials": result.get("credentials")})
@manager.route("/box/oauth/web/start", methods=["POST"]) # noqa: F821
@login_required
async def start_box_web_oauth():
req = await get_request_json()
client_id = req.get("client_id")
client_secret = req.get("client_secret")
redirect_uri = req.get("redirect_uri", BOX_WEB_OAUTH_REDIRECT_URI)
if not client_id or not client_secret:
return get_json_result(code=RetCode.ARGUMENT_ERROR, message="Box client_id and client_secret are required.")
flow_id = str(uuid.uuid4())
box_auth = BoxOAuth(
OAuthConfig(
client_id=client_id,
client_secret=client_secret,
)
)
auth_url = box_auth.get_authorize_url(
options=GetAuthorizeUrlOptions(
redirect_uri=redirect_uri,
state=flow_id,
)
)
cache_payload = {
"user_id": current_user.id,
"auth_url": auth_url,
"client_id": client_id,
"client_secret": client_secret,
"created_at": int(time.time()),
}
REDIS_CONN.set_obj(_web_state_cache_key(flow_id, "box"), cache_payload, WEB_FLOW_TTL_SECS)
return get_json_result(
data = {
"flow_id": flow_id,
"authorization_url": auth_url,
"expires_in": WEB_FLOW_TTL_SECS,}
)
@manager.route("/box/oauth/web/callback", methods=["GET"]) # noqa: F821
async def box_web_oauth_callback():
flow_id = request.args.get("state")
if not flow_id:
return await _render_web_oauth_popup("", False, "Missing OAuth parameters.", "box")
code = request.args.get("code")
if not code:
return await _render_web_oauth_popup(flow_id, False, "Missing authorization code from Box.", "box")
cache_payload = json.loads(REDIS_CONN.get(_web_state_cache_key(flow_id, "box")))
if not cache_payload:
return get_json_result(code=RetCode.ARGUMENT_ERROR, message="Box OAuth session expired or invalid.")
error = request.args.get("error")
error_description = request.args.get("error_description") or error
if error:
REDIS_CONN.delete(_web_state_cache_key(flow_id, "box"))
return await _render_web_oauth_popup(flow_id, False, error_description or "Authorization failed.", "box")
auth = BoxOAuth(
OAuthConfig(
client_id=cache_payload.get("client_id"),
client_secret=cache_payload.get("client_secret"),
)
)
auth.get_tokens_authorization_code_grant(code)
token = auth.retrieve_token()
result_payload = {
"user_id": cache_payload.get("user_id"),
"client_id": cache_payload.get("client_id"),
"client_secret": cache_payload.get("client_secret"),
"access_token": token.access_token,
"refresh_token": token.refresh_token,
}
REDIS_CONN.set_obj(_web_result_cache_key(flow_id, "box"), result_payload, WEB_FLOW_TTL_SECS)
REDIS_CONN.delete(_web_state_cache_key(flow_id, "box"))
return await _render_web_oauth_popup(flow_id, True, "Authorization completed successfully.", "box")
@manager.route("/box/oauth/web/result", methods=["POST"]) # noqa: F821
@login_required
@validate_request("flow_id")
async def poll_box_web_result():
req = await get_request_json()
flow_id = req.get("flow_id")
cache_blob = REDIS_CONN.get(_web_result_cache_key(flow_id, "box"))
if not cache_blob:
return get_json_result(code=RetCode.RUNNING, message="Authorization is still pending.")
cache_raw = json.loads(cache_blob)
if cache_raw.get("user_id") != current_user.id:
return get_json_result(code=RetCode.PERMISSION_ERROR, message="You are not allowed to access this authorization result.")
REDIS_CONN.delete(_web_result_cache_key(flow_id, "box"))
return get_json_result(data={"credentials": cache_raw})

View File

@ -13,7 +13,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
import logging
import asyncio
from api.db import UserTenantRole
from api.db.db_models import UserTenant
from api.db.services.user_service import UserTenantService, UserService
@ -24,7 +25,7 @@ from common.time_utils import delta_seconds
from api.utils.api_utils import get_data_error_result, get_json_result, get_request_json, server_error_response, validate_request
from api.utils.web_utils import send_invite_email
from common import settings
from api.apps import smtp_mail_server, login_required, current_user
from api.apps import login_required, current_user
@manager.route("/<tenant_id>/user/list", methods=["GET"]) # noqa: F821
@ -80,20 +81,24 @@ async def create(tenant_id):
role=UserTenantRole.INVITE,
status=StatusEnum.VALID.value)
if smtp_mail_server and settings.SMTP_CONF:
from threading import Thread
try:
user_name = ""
_, user = UserService.get_by_id(current_user.id)
if user:
user_name = user.nickname
Thread(
target=send_invite_email,
args=(invite_user_email, settings.MAIL_FRONTEND_URL, tenant_id, user_name or current_user.email),
daemon=True
).start()
asyncio.create_task(
send_invite_email(
to_email=invite_user_email,
invite_url=settings.MAIL_FRONTEND_URL,
tenant_id=tenant_id,
inviter=user_name or current_user.email
)
)
except Exception as e:
logging.exception(f"Failed to send invite email to {invite_user_email}: {e}")
return get_json_result(data=False, message="Failed to send invite email.", code=RetCode.SERVER_ERROR)
usr = invite_users[0].to_dict()
usr = {k: v for k, v in usr.items() if k in ["id", "avatar", "email", "nickname"]}

View File

@ -22,7 +22,7 @@ import secrets
import time
from datetime import datetime
from quart import redirect, request, session, make_response
from quart import redirect, request, session
from werkzeug.security import check_password_hash, generate_password_hash
from api.apps.auth import get_auth_client
@ -45,7 +45,7 @@ from api.utils.api_utils import (
)
from api.utils.crypt import decrypt
from rag.utils.redis_conn import REDIS_CONN
from api.apps import smtp_mail_server, login_required, current_user, login_user, logout_user
from api.apps import login_required, current_user, login_user, logout_user
from api.utils.web_utils import (
send_email_html,
OTP_LENGTH,
@ -868,9 +868,12 @@ async def forget_get_captcha():
from captcha.image import ImageCaptcha
image = ImageCaptcha(width=300, height=120, font_sizes=[50, 60, 70])
img_bytes = image.generate(captcha_text).read()
response = await make_response(img_bytes)
response.headers.set("Content-Type", "image/JPEG")
return response
import base64
base64_img = base64.b64encode(img_bytes).decode('utf-8')
data_uri = f"data:image/jpeg;base64,{base64_img}"
return get_json_result(data=data_uri)
@manager.route("/forget/otp", methods=["POST"]) # noqa: F821
@ -923,47 +926,58 @@ async def forget_send_otp():
ttl_min = OTP_TTL_SECONDS // 60
if not smtp_mail_server:
logging.warning("SMTP mail server not initialized; skip sending email.")
else:
try:
send_email_html(
subject="Your Password Reset Code",
to_email=email,
template_key="reset_code",
code=otp,
ttl_min=ttl_min,
)
except Exception:
return get_json_result(data=False, code=RetCode.SERVER_ERROR, message="failed to send email")
try:
await send_email_html(
subject="Your Password Reset Code",
to_email=email,
template_key="reset_code",
code=otp,
ttl_min=ttl_min,
)
except Exception as e:
logging.exception(e)
return get_json_result(data=False, code=RetCode.SERVER_ERROR, message="failed to send email")
return get_json_result(data=True, code=RetCode.SUCCESS, message="verification passed, email sent")
@manager.route("/forget", methods=["POST"]) # noqa: F821
async def forget():
"""
POST: Verify email + OTP and reset password, then log the user in.
Request JSON: { email, otp, new_password, confirm_new_password }
Deprecated single-step reset endpoint.
Use /forget/verify-otp then /forget/reset-password.
"""
return get_json_result(
data=False,
code=RetCode.NOT_EFFECTIVE,
message="Use /forget/verify-otp then /forget/reset-password",
)
def _verified_key(email: str) -> str:
return f"otp:verified:{email}"
@manager.route("/forget/verify-otp", methods=["POST"]) # noqa: F821
async def forget_verify_otp():
"""
Verify email + OTP only. On success:
- consume the OTP and attempt counters
- set a short-lived verified flag in Redis for the email
Request JSON: { email, otp }
"""
req = await get_request_json()
email = req.get("email") or ""
otp = (req.get("otp") or "").strip()
new_pwd = req.get("new_password")
new_pwd2 = req.get("confirm_new_password")
if not all([email, otp, new_pwd, new_pwd2]):
return get_json_result(data=False, code=RetCode.ARGUMENT_ERROR, message="email, otp and passwords are required")
# For reset, passwords are provided as-is (no decrypt needed)
if new_pwd != new_pwd2:
return get_json_result(data=False, code=RetCode.ARGUMENT_ERROR, message="passwords do not match")
if not all([email, otp]):
return get_json_result(data=False, code=RetCode.ARGUMENT_ERROR, message="email and otp are required")
users = UserService.query(email=email)
if not users:
return get_json_result(data=False, code=RetCode.DATA_ERROR, message="invalid email")
user = users[0]
# Verify OTP from Redis
k_code, k_attempts, k_last, k_lock = otp_keys(email)
if REDIS_CONN.get(k_lock):
@ -979,7 +993,6 @@ async def forget():
except Exception:
return get_json_result(data=False, code=RetCode.EXCEPTION_ERROR, message="otp storage corrupted")
# Case-insensitive verification: OTP generated uppercase
calc = hash_code(otp.upper(), salt)
if calc != stored_hash:
# bump attempts
@ -992,23 +1005,72 @@ async def forget():
REDIS_CONN.set(k_lock, int(time.time()), ATTEMPT_LOCK_SECONDS)
return get_json_result(data=False, code=RetCode.AUTHENTICATION_ERROR, message="expired otp")
# Success: consume OTP and reset password
# Success: consume OTP and attempts; mark verified
REDIS_CONN.delete(k_code)
REDIS_CONN.delete(k_attempts)
REDIS_CONN.delete(k_last)
REDIS_CONN.delete(k_lock)
# set verified flag with limited TTL, reuse OTP_TTL_SECONDS or smaller window
try:
REDIS_CONN.set(_verified_key(email), "1", OTP_TTL_SECONDS)
except Exception:
return get_json_result(data=False, code=RetCode.SERVER_ERROR, message="failed to set verification state")
return get_json_result(data=True, code=RetCode.SUCCESS, message="otp verified")
@manager.route("/forget/reset-password", methods=["POST"]) # noqa: F821
async def forget_reset_password():
"""
Reset password after successful OTP verification.
Requires: { email, new_password, confirm_new_password }
Steps:
- check verified flag in Redis
- update user password
- auto login
- clear verified flag
"""
req = await get_request_json()
email = req.get("email") or ""
new_pwd = req.get("new_password")
new_pwd2 = req.get("confirm_new_password")
if not all([email, new_pwd, new_pwd2]):
return get_json_result(data=False, code=RetCode.ARGUMENT_ERROR, message="email and passwords are required")
if new_pwd != new_pwd2:
return get_json_result(data=False, code=RetCode.ARGUMENT_ERROR, message="passwords do not match")
users = UserService.query(email=email)
if not users:
return get_json_result(data=False, code=RetCode.DATA_ERROR, message="invalid email")
user = users[0]
try:
UserService.update_user_password(user.id, new_pwd)
except Exception as e:
logging.exception(e)
return get_json_result(data=False, code=RetCode.EXCEPTION_ERROR, message="failed to reset password")
# Auto login (reuse login flow)
user.access_token = get_uuid()
login_user(user)
user.update_time = current_timestamp()
user.update_date = datetime_format(datetime.now())
user.save()
# login
try:
user.access_token = get_uuid()
login_user(user)
user.update_time = current_timestamp()
user.update_date = datetime_format(datetime.now())
user.save()
except Exception as e:
logging.exception(e)
return get_json_result(data=False, code=RetCode.EXCEPTION_ERROR, message="failed to login after reset")
# clear verified flag
try:
REDIS_CONN.delete(_verified_key(email))
except Exception:
pass
msg = "Password reset successful. Logged in."
return await construct_response(data=user.to_json(), auth=user.get_id(), message=msg)