dashboard_callbacks.py

import logging
import re
from typing import Dict
from html import escape
from uuid import uuid4
from bson.objectid import ObjectId
import datetime
from random import randint
from telegram import (
    InlineKeyboardButton,
    InlineKeyboardMarkup,
    Update,
    constants,
    InlineQueryResultArticle,
    InputTextMessageContent,
)
from telegram.ext import (
    ApplicationHandlerStop,
    ConversationHandler,
    ContextTypes,
)
from db_connection import uc

logging.basicConfig(
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    level=logging.INFO,
)
# for monitoring activities
logger = logging.getLogger(__name__)


# define states
# conv_dashboard state variables
# Pages
DASH_PAGE, FAV_PAGE = range(2)
# Callback data
DASH, FAV, SEARCH_CRED, ADD_CRED = range(4)
# FAV callback data
ONE, TWO, THREE, FOUR, FIVE, SIX = range(6)


def format_btns(fav_len, favs):
    fav_len = fav_len
    keyboard = []
    if fav_len == 1:
        keyboard.append(
            [
                InlineKeyboardButton(
                    f"{favs[0]}",
                    switch_inline_query_current_chat=f"🔎: {favs[0]}",
                ),
            ],
        )

    elif (fav_len % 2) == 1:
        count = 0
        for i in range(int((fav_len - 1) / 2)):
            keyboard.append(
                [
                    InlineKeyboardButton(
                        f"{favs[count]}",
                        switch_inline_query_current_chat=f"🔎: {favs[count]}",
                    ),
                    InlineKeyboardButton(
                        f"{favs[count+1]}",
                        switch_inline_query_current_chat=f"🔎: {favs[count+1]}",
                    ),
                ],
            )
            count += 2

        keyboard.append(
            [
                InlineKeyboardButton(
                    f"{favs[count]}",
                    switch_inline_query_current_chat=f"🔎: {favs[count]}",
                ),
            ],
        )

    elif (fav_len % 2) == 0:
        count = 0
        for i in range(int(fav_len / 2)):
            keyboard.append(
                [
                    InlineKeyboardButton(
                        f"{favs[count]}",
                        switch_inline_query_current_chat=f"🔎: {favs[count]}",
                    ),
                    InlineKeyboardButton(
                        f"{favs[count+1]}",
                        switch_inline_query_current_chat=f"🔎: {favs[count+1]}",
                    ),
                ],
            )
            count += 2

    return keyboard


def format_cred_query(cred_data: Dict[str, str]) -> str:
    """Hepler function for formatting the credential data"""
    text = f"""
<b>Name</b>: {cred_data["name"]}
<b>Username</b>: {cred_data["username"]}
<b>Password</b>: <code>{escape(cred_data["password"])}</code>
<b>URI</b>: {cred_data["uri"]}
<b>Tags</b>: {" ".join(cred_data["tags"])}
<b>Note</b>: {cred_data["note"]}
<b>Date</b>: {cred_data["date"]}
"""
    return text


async def dash(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    query = update.callback_query
    if not query:
        user = update.message.from_user
        logger.info(f"{user.first_name} has accessed dashbooard via command")

        doc = uc.find_one(
            {"user_id": f"{user.id}"},
            {"_id": 0, "credentials": 1, "favourites": 1},
        )
        context.user_data["credentials"] = doc["credentials"]
        context.user_data["favourites"] = doc["favourites"]

        keyboard = [
            [
                InlineKeyboardButton("⭐ FAV", callback_data="fav"),
                InlineKeyboardButton(
                    "🔎 SEARCH",
                    switch_inline_query_current_chat=str("🔎: "),
                ),
            ],
            [
                InlineKeyboardButton(
                    "🔑 GENERATE PASSWORD",
                    switch_inline_query_current_chat=str(randint(9, 25)),
                ),
            ],
            [
                InlineKeyboardButton(
                    "🔐 ADD CREDENTIAL",
                    callback_data="add_cred",
                ),
            ],
        ]
        reply_markup = InlineKeyboardMarkup(keyboard)
        await update.message.reply_text(
            f"""•══༻◉••Dashboard••◉༺══•
<b>Credentials</b>: {len(context.user_data["credentials"])}
<b>Favourites</b>: {len(context.user_data["favourites"])}
<b>Notes</b>: Coming Soon.
<b>Tags</b>: Coming Soon.
<b>Trash</b>: Coming Soon.
""",
            reply_markup=reply_markup,
            parse_mode=constants.ParseMode.HTML,
        )
        raise ApplicationHandlerStop

    else:
        user = query.from_user
        logger.info(f"{user.first_name} has accessed dashbooard via button")

        doc = uc.find_one(
            {"user_id": f"{user.id}"},
            {"_id": 0, "credentials": 1, "favourites": 1},
        )
        context.user_data["credentials"] = doc["credentials"]
        context.user_data["favourites"] = doc["favourites"]

        keyboard = [
            [
                InlineKeyboardButton("⭐ FAV", callback_data="fav"),
                InlineKeyboardButton(
                    "🔎 SEARCH",
                    switch_inline_query_current_chat=str("🔎: "),
                ),
            ],
            [
                InlineKeyboardButton(
                    "🔑 GENERATE PASSWORD",
                    switch_inline_query_current_chat=str(randint(9, 25)),
                ),
            ],
            [
                InlineKeyboardButton(
                    "🔐 ADD CREDENTIAL",
                    callback_data="add_cred",
                ),
            ],
        ]
        reply_markup = InlineKeyboardMarkup(keyboard)
        await query.edit_message_text(
            f"""•══༻◉••Dashboard••◉༺══•
<b>Credentials</b>: {len(context.user_data["credentials"])}
<b>Favourites</b>: {len(context.user_data["favourites"])}
<b>Notes</b>: Coming Soon.
<b>Tags</b>: Coming Soon.
<b>Trash</b>: Coming Soon.
""",
            reply_markup=reply_markup,
            parse_mode=constants.ParseMode.HTML,
        )
        raise ApplicationHandlerStop


async def fav(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    """Favourite credentials"""
    query = update.callback_query
    user = query.from_user
    logger.info(f"{user.first_name} is accessing favorites")

    doc = uc.find_one(
        {"user_id": f"{user.id}"},
        {"favourites": 1},
    )
    context.user_data["favourites"] = doc["favourites"]
    favs = context.user_data["favourites"]

    await query.answer()
    fav_len = len(favs)
    if fav_len == 0:
        keyboard = [
            [
                InlineKeyboardButton(
                    "<< BACK",
                    callback_data="dash",
                ),
            ]
        ]
        reply_markup = InlineKeyboardMarkup(keyboard)
        await query.edit_message_text(
            "Sorry you do not have any favourites:",
            reply_markup=reply_markup,
        )
        return ConversationHandler.END
    else:
        keyboard = format_btns(fav_len, favs)
        keyboard.append(
            [
                InlineKeyboardButton(
                    "<< BACK",
                    callback_data="dash",
                ),
            ],
        )
        reply_markup = InlineKeyboardMarkup(keyboard)
        await query.edit_message_text(
            "Here are your favourites:",
            reply_markup=reply_markup,
        )
        return ConversationHandler.END


# async def fav_query(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
#     query = update.callback_query
#     user = query.from_user
#     context.user_data["collection"] = uc.find_one(
#         {"user_id": f"{user.id}"}
#     )
#     creds = context.user_data.collection["credentials"]
#     favs = context.user_data.collection["favourites"]
#     fav_len = len(favs)

#     logger.info(f"{user.first_name} has accessed dashboard-fav from query")

#     await query.answer()
#     keys = [
#         [
#             InlineKeyboardButton(f"{favs[0]}", callback_data=str(ONE)),
#             InlineKeyboardButton(f"{favs[1]}", callback_data=str(TWO)),
#         ],
#         [
#             InlineKeyboardButton(f"{favs[2]}", callback_data=str(THREE)),
#             InlineKeyboardButton(f"{favs[3]}", callback_data=str(FOUR)),
#         ],
#         [
#             InlineKeyboardButton(f"{favs[4]}", callback_data=str(FIVE)),
#             InlineKeyboardButton(f"{favs[5]}", callback_data=str(SIX)),
#         ],
#     ]
#     keyboard = format_btns(fav_len, keys)
#     keyboard.append(
#         [
#             InlineKeyboardButton(
#                 "<< BACK",
#                 callback_data=str(DASH),
#             ),
#         ],
#     )
#     reply_markup = InlineKeyboardMarkup(keyboard)
#     await query.edit_message_text(
#         """Here are your favourites:
# """,
#         reply_markup=reply_markup,
#     )
#     return FAV_PAGE


async def done_query(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    """Done with credential"""
    query = update.callback_query
    user = query.from_user

    logger.info(f"{user.first_name} is done with credential")

    await query.answer()
    keyboard = [
        [
            InlineKeyboardButton(
                "🔐 Add Credential",
                callback_data="add_cred",
            ),
            InlineKeyboardButton(
                "🔎 Search",
                switch_inline_query_current_chat="🔎: ",
            ),
        ],
        [
            InlineKeyboardButton(
                "Dashboard",
                callback_data="dash",
            ),
        ],
    ]

    # await update.message.delete()
    await query.edit_message_text(
        text="All done!",
        reply_markup=InlineKeyboardMarkup(keyboard),
    )
    context.user_data.clear()
    return ConversationHandler.END


# get_cred
async def get_cred(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    """Search for credential"""
    keyboard = [
        [
            InlineKeyboardButton(
                "🔎 Search",
                switch_inline_query_current_chat="🔎: ",
            ),
        ],
    ]

    await update.message.delete()
    await update.message.reply_text(
        text="click 👇 to search.",
        reply_markup=InlineKeyboardMarkup(keyboard),
    )
    raise ApplicationHandlerStop


# async def chosen_cred(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
#     chosen_result = update.chosen_inline_result
#     chosen_id = chosen_result.inline_message_id
#     user = chosen_result.from_user

#     logger.info(f"{chosen_id}")

#     credentials = context.user_data["credentials"]
#     print(f"These are the credentials:\n{credentials}\n")

#     cred_result = {}
#     for i in credentials:
#         if str(i["_id"]) == chosen_id:
#             cred_result = i
#             break

#     logger.info(f"\n\n{cred_result}\n\n")

#     keyboard = [
#         [
#             InlineKeyboardButton(
#                 "EDIT",
#                 callback_data="edit_cred",
#             ),
#             InlineKeyboardButton(
#                 "DELETE",
#                 callback_data="delete_cred",
#             ),
#         ],
#         [
#             InlineKeyboardButton(
#                 "ADD TO FAV",
#                 callback_data="fav_cred",
#             ),
#         ],
#         [
#             InlineKeyboardButton(
#                 "DONE",
#                 callback_data="done",
#             ),
#         ],
#     ]
#     await user.send_message(
#         format_cred_query(cred_result),
#         parse_mode=constants.ParseMode.HTML,
#         reply_markup=InlineKeyboardMarkup(keyboard),
#     )
#     return ConversationHandler.END


async def inline_credential(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    query = update.inline_query.query
    query_user_id = update.inline_query.from_user.id

    if not re.search("^🔎: .*$", query):
        return

    query = re.split(" ", query, 1)[1]
    doc = uc.find_one(
        {"user_id": f"{query_user_id}"}, {"_id": 0, "credentials": 1, "favourites": 1}
    )
    context.user_data["credentials"] = doc["credentials"]
    context.user_data["favourites"] = doc["favourites"]

    credentials = context.user_data["credentials"]
    print(f"These are the credentials:\n{credentials}\n")
    favs = context.user_data["favourites"]
    print(f"These are the credentials:\n{favs}\n")

    cred_results = []
    for i in credentials:
        if re.search(f"^{query}.*$", i["name"], re.IGNORECASE):
            cred_results.append(i)

    logger.info(f"\n\n{cred_results}\n\n")

    results = []

    for i in cred_results:
        if i["name"] in favs:
            keyboard_fav = [
                [
                    InlineKeyboardButton(
                        "EDIT",
                        callback_data=f"edit_{str(i['_id'])}",
                    ),
                    InlineKeyboardButton(
                        "DELETE",
                        callback_data=f"del_{str(i['_id'])}",
                    ),
                ],
                [
                    InlineKeyboardButton(
                        "REMOVE FROM FAV",
                        callback_data="defav_cred",
                    ),
                ],
                [
                    InlineKeyboardButton(
                        "DONE",
                        callback_data="done",
                    ),
                ],
            ]
            results.append(
                InlineQueryResultArticle(
                    id=str(i["_id"]),
                    title=f"{i['name']}",
                    description=f"{i['username']}",
                    input_message_content=InputTextMessageContent(
                        format_cred_query(i),
                        parse_mode=constants.ParseMode.HTML,
                    ),
                    reply_markup=InlineKeyboardMarkup(keyboard_fav),
                ),
            )
        else:
            keyboard = [
                [
                    InlineKeyboardButton(
                        "EDIT",
                        callback_data=f"edit_{str(i['_id'])}",
                    ),
                    InlineKeyboardButton(
                        "DELETE",
                        callback_data=f"del_{str(i['_id'])}",
                    ),
                ],
                [
                    InlineKeyboardButton(
                        "ADD TO FAV",
                        callback_data="fav_cred",
                    ),
                ],
                [
                    InlineKeyboardButton(
                        "DONE",
                        callback_data="done",
                    ),
                ],
            ]
            results.append(
                InlineQueryResultArticle(
                    id=str(i["_id"]),
                    title=f"{i['name']}",
                    description=f"{i['username']}",
                    input_message_content=InputTextMessageContent(
                        format_cred_query(i),
                        parse_mode=constants.ParseMode.HTML,
                    ),
                    reply_markup=InlineKeyboardMarkup(keyboard),
                ),
            )

    await update.inline_query.answer(
        results,
        cache_time=0,
        is_personal=True,
    )