From b355ed5c0784181702d6032e7771bba74abd8081 Mon Sep 17 00:00:00 2001
From: BarsTiger <kit.barsic@gmail.com>
Date: Fri, 12 May 2023 17:36:09 +0300
Subject: [PATCH] Loading all sessions from folder

---
 .../callbacks/create_session/callbacks.py     |  7 +---
 .../menu/callbacks/load_sessions/__init__.py  |  1 +
 .../menu/callbacks/load_sessions/callbacks.py | 37 +++++++++++++++++++
 modules/menu/items/manage/fields.py           |  4 +-
 4 files changed, 42 insertions(+), 7 deletions(-)
 create mode 100644 modules/menu/callbacks/load_sessions/__init__.py
 create mode 100644 modules/menu/callbacks/load_sessions/callbacks.py

diff --git a/modules/menu/callbacks/create_session/callbacks.py b/modules/menu/callbacks/create_session/callbacks.py
index cb20db6..2884a0d 100644
--- a/modules/menu/callbacks/create_session/callbacks.py
+++ b/modules/menu/callbacks/create_session/callbacks.py
@@ -16,7 +16,7 @@ async def create_session_callback():
         await (client := GeneratedClient(name=session_name)).start()
         sessions[session_name] = SessionConfig(
             id=client.me.id,
-            phone=client.phone_number,
+            phone=client.me.phone_number,
             profile_name=client.me.first_name + (f' {client.me.last_name}' if client.me.last_name else ''),
             username=client.me.username
         ).json()
@@ -27,8 +27,3 @@ async def create_session_callback():
     except OperationalError:
         print('[red]Cannot create session file.[/] Try using different name...')
         input()
-
-    except Exception as e:
-        print(f'[red]Error:[/] {e}...')
-        input()
-        return
diff --git a/modules/menu/callbacks/load_sessions/__init__.py b/modules/menu/callbacks/load_sessions/__init__.py
new file mode 100644
index 0000000..cd7351c
--- /dev/null
+++ b/modules/menu/callbacks/load_sessions/__init__.py
@@ -0,0 +1 @@
+from .callbacks import load_sessions_callback
diff --git a/modules/menu/callbacks/load_sessions/callbacks.py b/modules/menu/callbacks/load_sessions/callbacks.py
new file mode 100644
index 0000000..9fe5305
--- /dev/null
+++ b/modules/menu/callbacks/load_sessions/callbacks.py
@@ -0,0 +1,37 @@
+from modules.decorators.callback import async_callback
+from modules.config import sessions
+from modules.client import GeneratedClient
+from modules.config.models import SessionConfig
+from rich import print
+import os
+
+
+@async_callback
+async def load_sessions_callback():
+    if input('Do you really want to load all .session files that are not in database? '
+             'It may cause bans of your accounts, because logging to many sessions from one IP at once may seem'
+             'suspicious for Telegram (y/N) ') not in ['y', 'Y']:
+        return
+    print('Loading...')
+    loaded = 0
+    for file in os.listdir():
+        if not file.endswith('.session'):
+            continue
+        file = file.replace('.session', '')
+        if file not in sessions.keys():
+            try:
+                await (client := GeneratedClient(name=file)).start()
+                sessions[file] = SessionConfig(
+                    id=client.me.id,
+                    phone=client.me.phone_number,
+                    profile_name=client.me.first_name + (f' {client.me.last_name}' if client.me.last_name else ''),
+                    username=client.me.username
+                ).json()
+                print(f'[green]Created[/] session {file}...')
+                await client.stop()
+                loaded += 1
+            except Exception as e:
+                print(f'[red]Error[/] loading {file} ({e}), skipping...')
+
+    print(f'[green]Created[/] {loaded} records...')
+    input()
diff --git a/modules/menu/items/manage/fields.py b/modules/menu/items/manage/fields.py
index 1ac6084..3d8b330 100644
--- a/modules/menu/items/manage/fields.py
+++ b/modules/menu/items/manage/fields.py
@@ -1,6 +1,7 @@
 from cursesmenu import CursesMenu
 from cursesmenu.items import FunctionItem
 from ...callbacks.overkill import overkill_callback
+from ...callbacks.load_sessions import load_sessions_callback
 from ...dynamic import update
 from ...dynamic import delete
 
@@ -9,6 +10,7 @@ def get_items_list(menu: CursesMenu):
     _ = [
         update.get_submenu_item,
         delete.get_submenu_item,
-        lambda _: FunctionItem('Delete overkill sessions', function=overkill_callback)
+        lambda _: FunctionItem('Delete overkill sessions', function=overkill_callback),
+        lambda _: FunctionItem('Load new sessions in folder', function=load_sessions_callback)
     ]
     return [x(menu) for x in _]