Mega.nz-Bot/megadl/helpers/database.py at be10b356fd1d9d2eaec0e5f98735df0bcd36fd7b · partiallywritten/Mega.nz-Bot · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
# Copyright (c) 2021 - Present partiallywritten
# Author: https://github.com/partiallywritten
# Project: https://github.com/partiallywritten/Mega.nz-Bot
# Description: Database functions

from os import getenv

from pymongo import AsyncMongoClient, ReturnDocument


class CypherDB:
    # use slots for faster var lookups
    # usually this isnt needed, like at all
    # i cant justify myself for adding this either
    # but wait maybe i can?
    # i use find_one_and_update in .add(...) maximum efficiency so i guess its not that conflicting
    __slots__ = ("mongoc", "coll_users")

    def __init__(self) -> None:
        self.mongoc = AsyncMongoClient(getenv("MONGO_URI"))
        self.coll_users = self.mongoc["mega_nz"]["users"]

    # <<<<<<<<<< Active user functions >>>>>>>>>> #
    async def add(self, user_id: int):
        result = await self.coll_users.find_one_and_update(
            {"_id": user_id},
            {
                "$setOnInsert": {
                    "_id": user_id,
                    "email": "",
                    "password": "",
                    "status": {"banned": False, "reason": ""},
                    "total_downloads": 0,
                    "total_uploads": 0,
                    "proxy": "",
                }
            },
            projection={"_id": 0, "status": 1},
            upsert=True,
            return_document=ReturnDocument.AFTER,
        )
        return result["status"]

    async def plus_fl_count(
        self, user_id: int, downloads: int | None = None, uploads: int | None = None
    ):
        if downloads:
            await self.coll_users.update_one(
                {"_id": user_id},
                {"$inc": {"total_downloads": downloads}},
            )
        elif uploads:
            await self.coll_users.update_one(
                {"_id": user_id},
                {"$inc": {"total_uploads": uploads}},
            )

    async def delete(self, user_id: int):
        await self.coll_users.delete_one({"_id": user_id})

    async def is_there(self, user_id: int, use_acc: bool = False):
        """
        Returns:
            {
                "email": "",
                "password": "",
                "status": {"banned": False, "reason": ""},
                "total_downloads": 0,
                "total_uploads": 0,
                "proxy": ""
            }
        """
        uid = {"_id": user_id}
        docu = await self.coll_users.find_one(uid, {"_id": 0})
        if not docu:
            return None
        if use_acc:
            return docu if not "" in (docu["email"], docu["password"]) else None
        else:
            return docu

    # <<<<<<<<<< Banned user functions >>>>>>>>>> #
    async def ban_user(self, user_id: int, reason: str = "Got banned :("):
        await self.coll_users.update_one(
            {"_id": user_id},
            {"$set": {"status": {"banned": True, "reason": reason}}},
        )

    async def unban_user(self, user_id: int, reason: str = "Got unbanned :)"):
        await self.coll_users.update_one(
            {"_id": user_id},
            {"$set": {"status": {"banned": False, "reason": reason}}},
        )

    # <<<<<<<<<< Mega functions >>>>>>>>>> #
    async def mega_login(self, user_id: int, email: str, password: str):
        await self.coll_users.update_one(
            {"_id": user_id},
            {"$set": {"email": email, "password": password}},
        )

    async def mega_logout(self, user_id: int):
        await self.coll_users.update_one(
            {"_id": user_id},
            {"$set": {"email": "", "password": ""}},
        )

    # this isnt used anywhere as far as from what i can tell but im not gonna risk it by deleting anything either
    # fuck it
    async def how_many(self):
        return (user["user_id"] async for user in self.coll_users.find({}))

    # <<<<<<<<<< Proxy functions >>>>>>>>>> #
    async def update_proxy(self, user_id: int, proxy: str):
        await self.coll_users.update_one(
            {"_id": user_id},
            {"$set": {"proxy": proxy}},
        )

    async def get_proxy(self, user_id: int):
        return await self.coll_users.find_one(
            {"_id": user_id}, {"proxy": 1}
        )