From f5ceb8927f2e7f2a9e29d62c8e4cef876f917249 Mon Sep 17 00:00:00 2001
From: Kyomotoi <1172294279@qq.com>
Date: Sat, 6 Feb 2021 00:32:26 +0800
Subject: =?UTF-8?q?=F0=9F=8F=97=20=F0=9F=92=A9=20=E6=9B=B4=E6=94=B9?=
 =?UTF-8?q?=E9=A1=B9=E7=9B=AE=E7=BB=93=E6=9E=84=EF=BC=8C=E4=BF=AE=E5=A4=8D?=
 =?UTF-8?q?=E5=95=A5b=20BUG?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ATRI/service/limit.py | 78 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 78 insertions(+)
 create mode 100644 ATRI/service/limit.py

(limited to 'ATRI/service/limit.py')

diff --git a/ATRI/service/limit.py b/ATRI/service/limit.py
new file mode 100644
index 0000000..d5a8a04
--- /dev/null
+++ b/ATRI/service/limit.py
@@ -0,0 +1,78 @@
+import json
+import aiofiles
+
+from ATRI.exceptions import WriteError
+from . import SERVICE_DIR
+
+
+class Limit:
+    
+    file_name = "limit.service.json"
+    path = SERVICE_DIR / file_name
+    path.parent.mkdir(exist_ok=True, parents=True)
+    try:
+        data = json.loads(path.read_bytes())
+    except:
+        data = {}
+    
+    @classmethod
+    def _filling_service(cls, service, group: int = None):
+        # Determine whether the service exists in global variables.
+        if service not in cls.data["global"]:
+            cls.data["global"][service] = True
+        # Similary, for group.
+        if service not in cls.data[group]:
+            cls.data[group][service] = True
+        
+    
+    @classmethod
+    def get_service(cls) -> dict:
+        return cls.data
+    
+    @classmethod
+    async def auth_service(
+        cls, service: str, group: int = None) -> bool:
+        cls.data.setdefault("global", {})
+        cls.data.setdefault(group, {})
+        cls._filling_service(service, group)
+        
+        try:
+            async with aiofiles.open(
+                cls.path, 'w', encoding='utf-8') as target:
+                await target.write(
+                    json.dumps(
+                        cls.data
+                    )
+                )
+        except WriteError:
+            raise WriteError("Writing file failed!")
+        
+        if cls.data["global"][service]:
+            return True if cls.data[group][service] else False
+        else:
+            return False
+    
+    @classmethod
+    async def control_service(
+        cls,
+        service: str,
+        status: bool,
+        group: int = None
+    ) -> None:
+        cls._filling_service(service, group)
+        
+        if group:
+            cls.data[group][service] = status
+        else:
+            cls.data["global"][service] = status
+        
+        try:
+            async with aiofiles.open(
+                cls.path, 'w', encoding='utf-8') as target:
+                await target.write(
+                    json.dumps(
+                        cls.data
+                    )
+                )
+        except WriteError:
+            raise WriteError("Writing file failed!")
-- 
cgit v1.2.3