(20241214) testing token keys instead of direct ids.

This commit is contained in:
2024-12-14 11:55:05 +05:30
parent 3a03dd4a61
commit c673c3511f
14 changed files with 164 additions and 60 deletions
+2 -2
View File
@@ -161,7 +161,7 @@ async def handle_gmail_callback() -> render_template:
# if they don't match, we reject the authorization:
auth_token = await current_app.mail_controller.get_token(
mongo_conn = current_app.data_mongo,
token_id = g.inbound_data["state"]
token_key = g.inbound_data["state"]
)
if (
(not auth_token) or
@@ -218,7 +218,7 @@ async def handle_gmail_callback() -> render_template:
db_conn = current_app.sql_writer,
mongo_conn = current_app.data_mongo,
session_token = g.inbound_headers.get("X-Session-Token"),
token_id = g.inbound_data["state"],
token_key = g.inbound_data["state"],
auth_token = auth_token
)
+8 -8
View File
@@ -170,12 +170,13 @@ async def request_oauth_authorization_url(
# Start by assuming failure:
auth_url = None
# ┏┓ ┏┳┓ ┓
# ┃┓┏┓┏┓┏┓┏┓┏┓╋┏┓ ┃ ┏┓┃┏┏┓┏┓ ┃
# ┗┛┗ ┛┗┗ ┛ ┗┻┗┗ ┻ ┗┛┛┗┗ ┛┗ ┻┗┻
# ┏┓ ┏┳┓ ┓ ┓┏
# ┃┓┏┓┏┓┏┓┏┓┏┓╋┏┓ ┃ ┏┓┃┏┏┓┏┓ ┃┫ ┏┓┓┏
# ┗┛┗ ┛┗┗ ┛ ┗┻┗┗ ┻ ┗┛┛┗┗ ┛┗ ┛┗┛┗ ┗┫
# ┛
# Make a user identifier from the session info:
token_id = await current_app.mail_controller.get_token_id(
token_key = await current_app.mail_controller.get_token_key(
db_conn = current_app.sql_writer,
mongo_conn = current_app.data_mongo,
auth_token = CoreAuthTokenModel(
@@ -189,13 +190,12 @@ async def request_oauth_authorization_url(
),
session_token = inbound_headers["X-Session-Token"]
)
if token_id is None:
if token_key is None:
return ResponseModel(
status_code = StatusCodes.FAILED,
http_code = HttpCodes.INTERNAL_SERVER_ERROR,
message = "failed to generate token id"
message = "failed to generate token key"
)
token_id = str(token_id)
# ┏┓ ┏┓┳┳┓ •┓
# ┣ ┏┓┏┓ ┃┓┃┃┃┏┓┓┃
@@ -206,7 +206,7 @@ async def request_oauth_authorization_url(
# Get the authorization URL:
auth_url = await current_app.gmail_client.get_authorization_url(
scopes = SCOPES_GMAIL_MAIL_MANAGEMENT,
state = token_id,
state = str(token_key),
access_type = "offline",
approval_prompt = "force",
include_granted_scopes = "true",
+29 -2
View File
@@ -156,17 +156,44 @@ async def get_one_mail(
:return: A standard response structure.
"""
# ┏┓ ┓ ┏┓┓ ┓
# ┣┫┓┏╋┣┓ ┃ ┣┓┏┓┏┃┏
# ┛┗┗┻┗┛┗ ┗┛┛┗┗ ┗┛┗
# If the session token is invalid/expired:
if kwargs.get("session_info") is None:
return ResponseModel(
status_code = StatusCodes.FAILED,
http_code = HttpCodes.UNAUTHORIZED
http_code = HttpCodes.UNAUTHORIZED,
messge = "invalid session"
)
# ┏┓ ┓ ┏┳┓ ┓
# ┣ ┏┓╋┏┣┓ ┃ ┏┓┃┏┏┓┏┓┏
# ┻ ┗ ┗┗┛┗ ┻ ┗┛┛┗┗ ┛┗┛
# We first load the authorization tokens:
auth_token = await current_app.mail_controller.get_token(
mongo_conn = current_app.data_mongo,
token_key = inbound_data.tokenKey,
)
# If we failed to load the authorization tokens:
if not auth_token:
return ResponseModel(
status_code = StatusCodes.FAILED,
http_code = HttpCodes.UNAUTHORIZED,
message = f"no such token key '{inbound_data.tokenKey}'"
)
# ┏┓ ┓ ┳┳┓ •┓
# ┣ ┏┓╋┏┣┓ ┃┃┃┏┓┓┃
# ┻ ┗ ┗┗┛┗ ┛ ┗┗┻┗┗
# Get the mail:
message = await current_app.mail_controller.get_one_mail(
mongo_conn = current_app.data_mongo,
token_id = inbound_data.tokenId,
token_id = auth_token.authTokenId,
message_id = inbound_data.messageId
)
+19 -7
View File
@@ -163,17 +163,29 @@ async def list_mails(
:return: A standard response structure.
"""
# ┏┓ ┓ ┏┓┓ ┓
# ┣┫┓┏╋┣┓ ┃ ┣┓┏┓┏┃┏
# ┛┗┗┻┗┛┗ ┗┛┛┗┗ ┗┛┗
# If the session token is invalid/expired:
if await token_check.is_not_authorized(
mongo_conn = current_app.data_mongo,
user_info = kwargs.get("session_info"),
token_ids = inbound_data.tokenIds
): return ResponseModel(
if kwargs.get("session_info") is None:
return ResponseModel(
status_code = StatusCodes.FAILED,
http_code = HttpCodes.UNAUTHORIZED,
message = "user not authorized to use this token"
messge = "invalid session"
)
# ┏┓ ┓• ┳┳┓ •┓
# ┣ ┏┓┃┓┏╋ ┃┃┃┏┓┓┃┏
# ┗┛┛┗┗┗┛┗ ┛ ┗┗┻┗┗┛
# Get the token ids from the token keys:
auth_tokens = await current_app.mail_controller.get_tokens(
mongo_conn = current_app.data_mongo,
token_keys = inbound_data.tokenKeys
)
token_ids = [t.authTokenId for t in auth_tokens]
# Build the additional filter:
additional_filter = {}
if inbound_data.tags: additional_filter["tags"] = {"$in": inbound_data.tags}
@@ -182,7 +194,7 @@ async def list_mails(
# Get the mails:
mails_list = await current_app.mail_controller.list_mails(
mongo_conn = current_app.data_mongo,
token_ids = inbound_data.tokenIds,
token_ids = token_ids,
limit = inbound_data.count,
skip = inbound_data.fromCount,
additional_filter = additional_filter
+1 -1
View File
@@ -146,7 +146,7 @@ async def sync_mails(
db_conn = current_app.sql_writer,
mongo_conn = current_app.data_mongo,
user_info = user_info,
token_id = inbound_data.tokenId,
token_key = inbound_data.tokenKey,
llm = current_app.llm,
force_sync = inbound_data.forceSync,
start_date = inbound_data.startDate,
+17 -2
View File
@@ -156,6 +156,10 @@ async def update_mail_tags(
:return: A standard response structure.
"""
# ┏┓ ┓ ┏┓┓ ┓
# ┣┫┓┏╋┣┓ ┃ ┣┓┏┓┏┃┏
# ┛┗┗┻┗┛┗ ┗┛┛┗┗ ┗┛┗
# If the session token is invalid/expired:
if kwargs.get("session_info") is None:
return ResponseModel(
@@ -163,10 +167,21 @@ async def update_mail_tags(
http_code = HttpCodes.UNAUTHORIZED
)
# Get the mail:
# ┳┳ ┓ ┳┳┓ •┓
# ┃┃┏┓┏┫┏┓╋┏┓ ┃┃┃┏┓┓┃
# ┗┛┣┛┗┻┗┻┗┗ ┛ ┗┗┻┗┗
# ┛
# get the token id from the token key:
auth_token = await current_app.mail_controller.get_token(
mongo_conn = current_app.data_mongo,
token_key = inbound_data.tokenKey
)
# Update the mail:
success = await current_app.mail_controller.update_tags(
mongo_conn = current_app.data_mongo,
token_id = inbound_data.tokenId,
token_id = auth_token.authTokenId,
message_id = inbound_data.messageId,
unset_tags = inbound_data.unsetTags,
set_tags = inbound_data.setTags
+26 -18
View File
@@ -238,7 +238,7 @@ class MailController:
# ┗┛┛┗┗┻┗┛┗┗━•┗┛
@staticmethod
async def get_token_id(
async def get_token_key(
db_conn: AsyncMySQL,
mongo_conn: AsyncMongo,
auth_token: CoreAuthTokenModel,
@@ -246,7 +246,7 @@ class MailController:
) -> ObjectId:
# Simply call the core model:
return await current_app.core_auth_token_controller.get_token_id(
return await current_app.core_auth_token_controller.get_token_key(
db_conn = db_conn,
mongo_conn = mongo_conn,
auth_token = auth_token,
@@ -260,7 +260,7 @@ class MailController:
async def set_token(
db_conn: AsyncMySQL,
mongo_conn: AsyncMongo,
token_id: ObjectId | str,
token_key: ObjectId | str,
auth_token: CoreAuthTokenModel,
session_token: str = None
) -> bool:
@@ -269,7 +269,7 @@ class MailController:
return await current_app.core_auth_token_controller.set_token(
db_conn = db_conn,
mongo_conn = mongo_conn,
token_id = token_id,
token_key = token_key,
auth_token = auth_token,
token_notes = {
"email": auth_token.token["email"],
@@ -282,13 +282,25 @@ class MailController:
@staticmethod
async def get_token(
mongo_conn: AsyncMongo,
token_id: ObjectId | str = None,
token_key: ObjectId | str = None,
) -> CoreAuthTokenModel | None:
# Simply call the core model:
return await current_app.core_auth_token_controller.get_token(
mongo_conn = mongo_conn,
token_id = token_id
token_key = token_key
)
@staticmethod
async def get_tokens(
mongo_conn: AsyncMongo,
token_keys: ObjectId | str = None,
) -> List[CoreAuthTokenModel] | None:
# Simply call the core model:
return await current_app.core_auth_token_controller.get_tokens(
mongo_conn = mongo_conn,
token_keys = token_keys
)
# ┏┓ ┳┳┓
@@ -303,7 +315,6 @@ class MailController:
self,
mongo_conn: AsyncMongo,
user_info: CoreUserInfoModel,
token_id: ObjectId,
auth_token: CoreAuthTokenModel,
mail_client: AsyncGMailClient,
google_tokens: GoogleAuthTokens,
@@ -320,11 +331,11 @@ class MailController:
if not force_sync:
mail_records = await current_app.core_message_controller.get_previews(
mongo_conn = mongo_conn,
token_ids = [ObjectId(token_id)],
token_ids = [ObjectId(auth_token.authTokenId)],
limit = 1,
skip = 0,
additional_filter = {
"tokenId": ObjectId(token_id),
"tokenId": ObjectId(auth_token.authTokenId),
"serviceType": auth_token.serviceType,
"client": auth_token.client,
"clientMessageId": message_id
@@ -357,7 +368,7 @@ class MailController:
mail_message = CoreMessageModel(
ts = client_response.data["ts"],
syncTs = date_time.get_current_utc_date_time(as_string = False),
tokenId = token_id,
tokenId = auth_token.authTokenId,
serviceType = auth_token.serviceType,
client = auth_token.client,
clientMessageId = message_id,
@@ -398,7 +409,6 @@ class MailController:
db_conn: AsyncMySQL,
mongo_conn: AsyncMongo,
user_info: CoreUserInfoModel,
token_id: ObjectId,
auth_token: CoreAuthTokenModel,
mail_client: AsyncGMailClient,
llm: LLMController = None,
@@ -427,7 +437,7 @@ class MailController:
await self.set_token(
db_conn = db_conn,
mongo_conn = mongo_conn,
token_id = token_id,
token_key = auth_token.key,
auth_token = auth_token,
session_token = session_token
)
@@ -454,7 +464,6 @@ class MailController:
self.__sync_one_gmail(
mongo_conn = mongo_conn,
user_info = user_info,
token_id = token_id,
auth_token = auth_token,
mail_client = mail_client,
google_tokens = google_tokens,
@@ -478,7 +487,7 @@ class MailController:
mongo_operations.append(
ReplaceOne(
filter = {
"tokenId": ObjectId(token_id),
"tokenId": ObjectId(auth_token.authTokenId),
"serviceType": auth_token.serviceType,
"client": auth_token.client,
"clientMessageId": result.mailMessage.clientMessageId
@@ -514,7 +523,7 @@ class MailController:
db_conn: AsyncMySQL,
mongo_conn: AsyncMongo,
user_info: CoreUserInfoModel,
token_id: ObjectId | str,
token_key: ObjectId | str,
llm: LLMController = None,
force_sync: bool = False,
start_date: datetime.datetime = None,
@@ -533,12 +542,12 @@ class MailController:
# We first load the authorization tokens:
auth_token = await self.get_token(
mongo_conn = mongo_conn,
token_id = token_id,
token_key = token_key,
)
# If we failed to load the authorization tokens:
if not auth_token:
sync_results.message = f"no such token id '{token_id}'"
sync_results.message = f"no such token key '{token_key}'"
return sync_results
# ┏┓ ┏┓┳┳┓ •┓
@@ -550,7 +559,6 @@ class MailController:
db_conn = db_conn,
mongo_conn = mongo_conn,
user_info = user_info,
token_id = token_id,
auth_token = auth_token,
mail_client = current_app.gmail_client,
llm = llm,
+14 -13
View File
@@ -99,7 +99,7 @@ class AuthTokenController(BaseModel):
# For MongoDB:
AUTH_COLLECTION = "_authTokens"
async def get_token_id(
async def get_token_key(
self,
db_conn: AsyncMySQL,
mongo_conn: AsyncMongo,
@@ -141,6 +141,7 @@ class AuthTokenController(BaseModel):
"syncFreq": auth_token.syncFreq
},
"$setOnInsert": {
"key": auth_token.key,
"serviceType": auth_token.serviceType,
"client": auth_token.client,
"authType": auth_token.authType,
@@ -173,7 +174,7 @@ class AuthTokenController(BaseModel):
"Auth Requested", # ............................................. 'p_last_action'
None, # ......................................................... 'p_display_name'
None, # ......................................................... 'p_display_picture'
str(mongo_json["_id"]), # ....................................... 'p_token_id'
str(auth_token.key), # .......................................... 'p_token_id'
json.to_string(python_data = token_notes, no_space = True), # ... 'p_notes'
auth_token.user.userId # ........................................ 'p_created_by'
),
@@ -181,13 +182,13 @@ class AuthTokenController(BaseModel):
)
# Done here:
return mongo_json["_id"] if mongo_json and db_json.get("status") == 1 else None
return auth_token.key if mongo_json and db_json.get("status") == 1 else None
async def set_token(
self,
db_conn: AsyncMySQL,
mongo_conn: AsyncMongo,
token_id: ObjectId | str,
token_key: ObjectId | str,
auth_token: CoreAuthTokenModel,
token_notes: dict,
session_token: str = None
@@ -199,7 +200,7 @@ class AuthTokenController(BaseModel):
ALSO.
:param db_conn: The database connection (MariaDB) to use to perform the action.
:param mongo_conn: The database connection (MongoDB) to use to perform the action.
:param token_id: The identifier granted by the 'get_token_id' method.
:param token_key: The identifier granted by the 'get_token_key' method.
:param auth_token: The actual auth/token data to be saved to the database.
:param token_notes: Any notes to feed into MariaDB with the token identifier.
:param session_token: The session token of the user who requested this service.
@@ -217,7 +218,7 @@ class AuthTokenController(BaseModel):
mongo_json = await mongo_conn.find_one_and_update(
collection = self.AUTH_COLLECTION,
filter = mongo_conn.dict_to_dot_notation({
"_id": ObjectId(token_id),
"key": ObjectId(token_key),
"clientUserId": auth_token.clientUserId
}),
update = [{
@@ -257,7 +258,7 @@ class AuthTokenController(BaseModel):
"Auth Granted", # ............................................... 'p_last_action'
auth_token.token.get("displayName"), # .......................... 'p_display_name'
auth_token.token.get("displayPictureUrl"), # .................... 'p_display_picture'
token_id, # ..................................................... 'p_token_id'
token_key, # .................................................... 'p_token_id'
json.to_string(python_data = token_notes, no_space = True), # ... 'p_notes'
auth_token.user.userId # ........................................ 'p_created_by'
),
@@ -271,13 +272,13 @@ class AuthTokenController(BaseModel):
async def get_token(
self,
mongo_conn: AsyncMongo,
token_id: ObjectId | str = None,
token_key: ObjectId | str = None,
) -> CoreAuthTokenModel | None:
"""
To retrieve stored tokens from the database. One token at a time.
:param mongo_conn: The database connection (MongoDB) to use to perform the action.
:param token_id: The identifier granted by the 'get_token_id' method.
:param token_key: The identifier granted by the 'get_token_key' method.
:return: The retrieved record that has the token, and information about the service and client if found, else
None when there is no matching record.
"""
@@ -285,7 +286,7 @@ class AuthTokenController(BaseModel):
# If there is some filtering possible, we fetch the token:
token = await mongo_conn.find_one(
collection = self.AUTH_COLLECTION,
filter = {"_id": ObjectId(token_id)},
filter = {"key": ObjectId(token_key)},
)
# Done here:
@@ -294,13 +295,13 @@ class AuthTokenController(BaseModel):
async def get_tokens(
self,
mongo_conn: AsyncMongo,
token_ids: List[ObjectId | str] = None,
token_keys: List[ObjectId | str] = None,
) -> List[CoreAuthTokenModel]:
"""
To retrieve stored tokens from the database. Multiple tokens at a time.
:param mongo_conn: The database connection (MongoDB) to use to perform the action.
:param token_ids: the identifiers granted by the 'get_token_id' method.
:param token_keys: the identifiers granted by the 'get_token_key' method.
:return: The retrieved record that has the token, and information about the service and client if found, else
None when there is no matching record.
"""
@@ -308,7 +309,7 @@ class AuthTokenController(BaseModel):
# If there is some filtering possible, we fetch the token:
tokens = await mongo_conn.find_many(
collection = self.AUTH_COLLECTION,
filter = {"_id": {"$in": [ObjectId(t) for t in token_ids]}}
filter = {"key": {"$in": [ObjectId(k) for k in token_keys]}}
)
# Done here:
+1 -1
View File
@@ -101,7 +101,7 @@ class MailGetRequestHeaders(BaseModel):
class MailGetRequestData(BaseModel):
tokenId: str = Field(
tokenKey: str = Field(
description = "the id of the token associated with the mail; needed for security",
frozen = True
)
+2 -2
View File
@@ -101,9 +101,9 @@ class MailListRequestHeaders(BaseModel):
class MailListRequestData(BaseModel):
tokenIds: str | List[str] = Field(
tokenKeys: str | List[str] = Field(
description = "the token identifier(s) that tell you which auth-tokens were used for fetching those messages",
frozen = True
frozen = True,
)
count: int = Field(
+1 -1
View File
@@ -104,7 +104,7 @@ class MailSyncRequestHeaders(BaseModel):
class MailSyncRequestData(BaseModel):
tokenId: str = Field(
tokenKey: str = Field(
description = "the account identifier (Mongo ObjectId) granted by 'MailOAuthModel.get_account_identifier'",
frozen = True
)
+1 -1
View File
@@ -101,7 +101,7 @@ class MailUpdateTagsRequestHeaders(BaseModel):
class MailUpdateTagsRequestData(BaseModel):
tokenId: str = Field(
tokenKey: str = Field(
description = "the id of the token associated with the mail; needed for security",
frozen = True
)
+7 -1
View File
@@ -83,12 +83,18 @@ import datetime
class CoreAuthTokenModel(BaseModel):
authTokenId: ObjectId = Field(
description = "the id of the document in mongodb that holds this information",
description = "the id of the document in mongodb that holds this information; hide from the ui layer",
frozen = True,
default = None,
alias = "_id"
)
key: ObjectId = Field(
description = "the expendable reference to this auth; expose this to the ui",
frozen = True,
default_factory = lambda: ObjectId()
)
serviceType: Literal["software", "email", "sms", "chat", "paymentGateway"] = Field(
description = "the kind of service this message was sent/received from",
frozen = True
+35
View File
@@ -0,0 +1,35 @@
import asyncio
from utils_v2.database.async_mongo_v2 import AsyncMongo
async def main():
# MongoDB connections:
data_mongo = AsyncMongo(
connection_string = r"mongodb://del.ditscentre.in:27017,wtt.ditscentre.in:27017,mum.arh.001.ditscentre.in:27017/admin?tls=true&tlsCAFile=%2Fetc%2Fssl%2Fcerts%2Fmongo_data_ca.pem&tlsCertificateKeyFile=%2Fetc%2Fssl%2Fcerts%2Fmongo_data_cert.pem&replicaSet=dits_mongod_rep&readPreference=primary&authMechanism=MONGODB-X509&authSource=%24external",
database_name = "converse",
max_connections = True,
debug = True
)
await data_mongo.connect()
while True:
updated = await data_mongo.find_one_and_update(
collection = "_authTokens",
filter = {
"$or": [
{"key": None},
{"key": {"$exists": False}}
]
},
update = {
"$set": {
"key": data_mongo.generate_id(as_str = False)
}
}
)
if not updated: break
print("UPDATED:", str(updated["_id"]))
asyncio.run(main())