feat: main feature is completed

This commit is contained in:
Hehesheng 2024-05-06 23:14:09 +08:00
parent 13b4ef6e66
commit 92850c1ba1
9 changed files with 452 additions and 79 deletions

2
.gitignore vendored
View File

@ -3,3 +3,5 @@ __pycache__
.vscode .vscode
*.session *.session
*.toml *.toml
*.db
*.service

View File

@ -1,9 +1,11 @@
import asyncio import asyncio
from typing import Union import json
from typing import Union, Optional
from telethon import TelegramClient, types from telethon import TelegramClient, types, hints
import configParse import configParse
import apiutils
class TgFileSystemClient(object): class TgFileSystemClient(object):
@ -12,12 +14,13 @@ class TgFileSystemClient(object):
session_name: str session_name: str
proxy_param: dict[str, any] proxy_param: dict[str, any]
client: TelegramClient client: TelegramClient
dialogs_cache: Optional[hints.TotalList] = None
me: Union[types.User, types.InputPeerUser] me: Union[types.User, types.InputPeerUser]
def __init__(self, param: configParse.TgToFileSystemParameter) -> None: def __init__(self, session_name: str, param: configParse.TgToFileSystemParameter) -> None:
self.api_id = param.tgApi.api_id self.api_id = param.tgApi.api_id
self.api_hash = param.tgApi.api_hash self.api_hash = param.tgApi.api_hash
self.session_name = param.base.name self.session_name = session_name
self.proxy_param = { self.proxy_param = {
'proxy_type': param.proxy.proxy_type, 'proxy_type': param.proxy.proxy_type,
'addr': param.proxy.addr, 'addr': param.proxy.addr,
@ -26,26 +29,108 @@ class TgFileSystemClient(object):
self.client = TelegramClient( self.client = TelegramClient(
self.session_name, self.api_id, self.api_hash, proxy=self.proxy_param) self.session_name, self.api_id, self.api_hash, proxy=self.proxy_param)
def __del__(self) -> None:
self.client.disconnect()
def __repr__(self) -> str: def __repr__(self) -> str:
if not self.client.is_connected: if not self.client.is_connected:
return f"client disconnected, session_name:{self.session_name}" return f"client disconnected, session_name:{self.session_name}"
return f"client connected, session_name:{self.session_name}, username:{self.me.username}, phone:{self.me.phone}, detail:{self.me.stringify()}" return f"client connected, session_name:{self.session_name}, username:{self.me.username}, phone:{self.me.phone}, detail:{self.me.stringify()}"
async def init_client(self): def _call_before_check(func):
def call_check_wrapper(self, *args, **kwargs):
if not self.is_valid():
raise RuntimeError("Client does not run.")
result = func(self, *args, **kwargs)
return result
return call_check_wrapper
def _acall_before_check(func):
async def call_check_wrapper(self, *args, **kwargs):
if not self.is_valid():
raise RuntimeError("Client does not run.")
result = await func(self, *args, **kwargs)
return result
return call_check_wrapper
@_call_before_check
def to_dict(self) -> dict:
return self.me.to_dict()
@_call_before_check
def to_json(self) -> str:
return self.me.to_json()
def is_valid(self) -> bool:
return self.client.is_connected() and self.me is not None
async def start(self) -> None:
if not self.client.is_connected():
await self.client.connect()
self.me = await self.client.get_me() self.me = await self.client.get_me()
if self.me is None:
raise RuntimeError(
f"The {self.session_name} Client Does Not Login")
async def stop(self) -> None:
await self.client.disconnect()
@_acall_before_check
async def get_message(self, chat_id: int, msg_id: int) -> types.Message:
msg = await self.client.get_messages(chat_id, ids=msg_id)
return msg
@_acall_before_check
async def get_dialogs(self, limit: int = 10, offset: int = 0, refresh: bool = False) -> hints.TotalList:
def _to_json(item) -> str:
return json.dumps({"id": item.id, "is_channel": item.is_channel,
"is_group": item.is_group, "is_user": item.is_user, "name": item.name, })
if self.dialogs_cache is not None and refresh is False:
return self.dialogs_cache[offset:offset+limit]
self.dialogs_cache = await self.client.get_dialogs()
for item in self.dialogs_cache:
item.to_json = _to_json
return self.dialogs_cache[offset:offset+limit]
async def _get_offset_msg_id(self, chat_id: int, offset: int) -> int:
if offset != 0:
begin = await self.client.get_messages(chat_id, limit=1)
if len(begin) == 0:
return hints.TotalList()
first_id = begin[0].id
offset = first_id + offset
return offset
@_acall_before_check
async def get_messages(self, chat_id: int, limit: int = 10, offset: int = 0) -> hints.TotalList:
offset = await self._get_offset_msg_id(chat_id, offset)
res_list = await self.client.get_messages(chat_id, limit=limit, offset_id=offset)
return res_list
@_acall_before_check
async def get_messages_by_search(self, chat_id: int, search_word: str, limit: int = 10, offset: int = 0, inner_search: bool = False) -> hints.TotalList:
offset = await self._get_offset_msg_id(chat_id, offset)
if inner_search:
res_list = await self.client.get_messages(chat_id, limit=limit, offset_id=offset, search=search_word)
return res_list
# search by myself
res_list = hints.TotalList()
async for msg in self.client.iter_messages(chat_id, offset_id=offset):
if msg.text.find(search_word) == -1 and apiutils.get_message_media_name(msg).find(search_word) == -1:
continue
res_list.append(msg)
if len(res_list) >= limit:
break
return res_list
def __enter__(self): def __enter__(self):
self.client.__enter__() raise NotImplemented
self.client.loop.run_until_complete(self.init_client())
def __exit__(self): def __exit__(self):
self.client.__exit__() raise NotImplemented
self.me = None
async def __aenter__(self): async def __aenter__(self):
await self.client.__enter__() await self.start()
await self.init_client()
async def __aexit__(self): async def __aexit__(self):
await self.client.__aexit__() await self.stop()

View File

@ -1,35 +1,63 @@
from typing import Any from typing import Any
import time
import hashlib
import os
from TgFileSystemClient import TgFileSystemClient from TgFileSystemClient import TgFileSystemClient
from UserManager import UserManager
import configParse
class TgFileSystemClientManager(object): class TgFileSystemClientManager(object):
MAX_MANAGE_CLIENTS: int = 10 MAX_MANAGE_CLIENTS: int = 10
clients: dict[int, TgFileSystemClient] param: configParse.TgToFileSystemParameter
clients: dict[str, TgFileSystemClient] = {}
def __init__(self) -> None: def __init__(self, param: configParse.TgToFileSystemParameter) -> None:
self.param = param
self.db = UserManager()
def __del__(self) -> None:
pass pass
def push_client(self, client: TgFileSystemClient) -> int: def check_client_session_exist(self, client_id: str) -> bool:
""" return os.path.isfile(client_id + '.session')
push client to manager.
Arguments def generate_client_id(self) -> str:
client return hashlib.md5(
(str(time.perf_counter()) + self.param.base.salt).encode('utf-8')).hexdigest()
Returns def create_client(self, client_id: str = None) -> TgFileSystemClient:
client id if client_id is None:
client_id = self.generate_client_id()
client = TgFileSystemClient(client_id, self.param)
return client
""" def register_client(self, client: TgFileSystemClient) -> bool:
self.clients[id(client)] = client self.clients[client.session_name] = client
return id(client) return True
def get_client(self, client_id: int) -> TgFileSystemClient: def deregister_client(self, client_id: str) -> bool:
self.clients.pop(client_id)
return True
def get_client(self, client_id: str) -> TgFileSystemClient:
client = self.clients.get(client_id) client = self.clients.get(client_id)
return client return client
async def get_client_force(self, client_id: str) -> TgFileSystemClient:
client = self.get_client(client_id)
if client is None:
if not self.check_client_session_exist(client_id):
raise RuntimeError("Client session does not found.")
client = self.create_client(client_id=client_id)
if not client.is_valid():
await client.start()
self.register_client(client)
return client
if __name__ == "__main__": if __name__ == "__main__":
import configParse import configParse
t: TgFileSystemClient = TgFileSystemClient(configParse.get_TgToFileSystemParameter()) # t: TgFileSystemClient = TgFileSystemClient(configParse.get_TgToFileSystemParameter())
print(f"{t.session_name=}") print(f"{t.session_name=}")

59
UserManager.py Normal file
View File

@ -0,0 +1,59 @@
import sqlite3
from pydantic import BaseModel
class UserUpdateParam(BaseModel):
client_id: str
username: str
phone: str
tg_user_id: int
last_login_time: int
class MessageUpdateParam(BaseModel):
tg_chat_id: int
tg_message_id: int
client_id: str
username: str
phone: str
tg_user_id: int
class UserManager(object):
def __init__(self) -> None:
self.con = sqlite3.connect("user.db")
self.cur = self.con.cursor()
if not self._table_has_been_inited():
self._first_runtime_run_once()
def __del__(self) -> None:
self.con.commit()
self.con.close()
def update_user(self) -> None:
raise NotImplemented
def update_message(self) -> None:
raise NotImplemented
def get_user_info() -> None:
raise NotImplemented
def _table_has_been_inited(self) -> bool:
res = self.cur.execute("SELECT name FROM sqlite_master")
return len(res.fetchall()) != 0
def _first_runtime_run_once(self) -> None:
if len(self.cur.execute("SELECT name FROM sqlite_master WHERE name='user'").fetchall()) == 0:
self.cur.execute(
"CREATE TABLE user(client_id, username, phone, tg_user_id, last_login_time)")
if len(self.cur.execute("SELECT name FROM sqlite_master WHERE name='message'").fetchall()) == 0:
self.cur.execute(
"CREATE TABLE message(tg_chat_id, tg_message_id, client_id, username, phone, tg_user_id, msg_ctx, msg_type)")
if __name__ == "__main__":
db = UserManager()
res = db.cur.execute("SELECT name FROM sqlite_master")
print(res.fetchall())

64
apiutils.py Normal file
View File

@ -0,0 +1,64 @@
import time
from fastapi import status, HTTPException
from telethon import types
from functools import wraps
import configParse
def get_range_header(range_header: str, file_size: int) -> tuple[int, int]:
def _invalid_range():
return HTTPException(
status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
detail=f"Invalid request range (Range:{range_header!r})",
)
try:
h = range_header.replace("bytes=", "").split("-")
start = int(h[0]) if h[0] != "" else 0
end = int(h[1]) if h[1] != "" else file_size - 1
except ValueError:
raise _invalid_range()
if start > end or start < 0 or end > file_size - 1:
raise _invalid_range()
return start, end
def get_message_media_name(msg: types.Message) -> str:
if msg.media is None or msg.media.document is None:
return ""
for attr in msg.media.document.attributes:
if isinstance(attr, types.DocumentAttributeFilename):
return attr.file_name
def timeit(func):
if configParse.get_TgToFileSystemParameter().base.timeit_enable:
@wraps(func)
def timeit_wrapper(*args, **kwargs):
start_time = time.perf_counter()
result = func(*args, **kwargs)
end_time = time.perf_counter()
total_time = end_time - start_time
print(
f'Function {func.__name__}{args} {kwargs} Took {total_time:.4f} seconds')
return result
return timeit_wrapper
return func
def atimeit(func):
if configParse.get_TgToFileSystemParameter().base.timeit_enable:
@wraps(func)
async def timeit_wrapper(*args, **kwargs):
start_time = time.perf_counter()
result = await func(*args, **kwargs)
end_time = time.perf_counter()
total_time = end_time - start_time
print(
f'AFunction {func.__name__}{args} {kwargs} Took {total_time:.4f} seconds')
return result
return timeit_wrapper
return func

14
config.toml.example Normal file
View File

@ -0,0 +1,14 @@
[base]
salt = "AnyTokenYouWanted"
port = 7777
timeit_enable = false
[tgApi]
api_id = int_app_id_from_tg
api_hash = "api_hash_from_tg"
[proxy]
enable = false
proxy_type = "socks5"
addr = "172.25.32.1"
port = 7890

View File

@ -4,8 +4,9 @@ from pydantic import BaseModel
class TgToFileSystemParameter(BaseModel): class TgToFileSystemParameter(BaseModel):
class BaseParameter(BaseModel): class BaseParameter(BaseModel):
name: str salt: str
port: int port: int
timeit_enable: bool
base: BaseParameter base: BaseParameter
class ApiParameter(BaseModel): class ApiParameter(BaseModel):

144
start.py
View File

@ -1,22 +1,29 @@
import asyncio import asyncio
import time
import json
import uvicorn import uvicorn
from fastapi import FastAPI from fastapi import FastAPI, status, Request
from fastapi import status
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response from fastapi.responses import Response, StreamingResponse
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from telethon import TelegramClient from telethon import types, hints
from pydantic import BaseModel
import configParse import configParse
import apiutils
from TgFileSystemClientManager import TgFileSystemClientManager
from TgFileSystemClient import TgFileSystemClient
clients_mgr: TgFileSystemClientManager = None
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
global clients_mgr
param = configParse.get_TgToFileSystemParameter() param = configParse.get_TgToFileSystemParameter()
loop = asyncio.get_event_loop() clients_mgr = TgFileSystemClientManager(param)
tg_client_task = loop.create_task(start_tg_client(param))
yield yield
asyncio.gather(*[tg_client_task])
app = FastAPI(lifespan=lifespan) app = FastAPI(lifespan=lifespan)
@ -28,47 +35,100 @@ app.add_middleware(
allow_headers=["*"], allow_headers=["*"],
) )
@app.post("/tg/{chat_id}/{message_id}")
async def get_test(chat_id: str, message_id: str): class TgToFileListRequestBody(BaseModel):
print(f"test: {chat_id=}, {message_id=}") token: str
return Response(status_code=status.HTTP_200_OK) search: str = ""
chat_id: int = 0
index: int = 0
length: int = 10
refresh: bool = False
inner: bool = False
async def start_tg_client(param: configParse.TgToFileSystemParameter): @app.post("/tg/api/v1/file/list")
api_id = param.tgApi.api_id @apiutils.atimeit
api_hash = param.tgApi.api_hash async def get_tg_file_list(body: TgToFileListRequestBody):
session_name = param.base.name try:
proxy_param = { res = hints.TotalList()
'proxy_type': param.proxy.proxy_type, res_type = "chat"
'addr': param.proxy.addr, client = await clients_mgr.get_client_force(body.token)
'port': param.proxy.port, res_dict = {}
} if param.proxy.enable else {} if body.chat_id == 0:
client = TelegramClient(session_name, api_id, api_hash, proxy=proxy_param) res = await client.get_dialogs(limit=body.length, offset=body.index, refresh=body.refresh)
res_dict = [{"id": item.id, "is_channel": item.is_channel,
"is_group": item.is_group, "is_user": item.is_user, "name": item.name, } for item in res]
elif body.search != "":
res = await client.get_messages_by_search(body.chat_id, search_word=body.search, limit=body.length, offset=body.index, inner_search=body.inner)
res_type = "msg"
res_dict = [json.loads(item.to_json()) for item in res]
else:
res = await client.get_messages(body.chat_id, limit=body.length, offset=body.index)
res_type = "msg"
res_dict = [json.loads(item.to_json()) for item in res]
async def tg_client_main(): response_dict = {
# Getting information about yourself "client": json.loads(client.to_json()),
me = await client.get_me() "type": res_type,
"list": res_dict,
}
return Response(json.dumps(response_dict), status_code=status.HTTP_200_OK)
except Exception as err:
print(f"{err=}")
return Response(f"{err=}", status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
# "me" is a user object. You can pretty-print
# any Telegram object with the "stringify" method:
print(me.stringify())
# When you print something, you see a representation of it. @app.get("/tg/api/v1/file/msg")
# You can access all attributes of Telegram objects with @apiutils.atimeit
# the dot operator. For example, to get the username: async def get_tg_file_media_stream(token: str, cid: int, mid: int, request: Request):
username = me.username async def get_msg_media_range_requests(client: TgFileSystemClient, msg: types.Message, start: int, end: int):
print(username) MAX_CHUNK_SIZE = 1024 * 1024
print(me.phone) pos = start
# You can print all the dialogs/conversations that you are part of: async for chunk in client.client.iter_download(msg, offset=pos, chunk_size=min(end + 1 - pos, MAX_CHUNK_SIZE)):
dialogs = await client.get_dialogs() pos = pos + len(chunk)
for dialog in dialogs: yield chunk.tobytes()
print(f"{dialog.name} has ID {dialog.id}") msg_id = mid
# async for dialog in client.iter_dialogs(): chat_id = cid
# print(dialog.name, 'has ID', dialog.id) headers = {
# "content-type": "video/mp4",
async with client: "accept-ranges": "bytes",
await tg_client_main() "content-encoding": "identity",
# "content-length": stream_file_size,
"access-control-expose-headers": (
"content-type, accept-ranges, content-length, "
"content-range, content-encoding"
),
}
range_header = request.headers.get("range")
try:
client = await clients_mgr.get_client_force(token)
msg = await client.get_message(chat_id, msg_id)
file_size = msg.media.document.size
start = 0
end = file_size - 1
status_code = status.HTTP_200_OK
mime_type = msg.media.document.mime_type
headers["content-type"] = mime_type
file_name = apiutils.get_message_media_name(msg)
if file_name == "":
maybe_file_type = mime_type.split("/")[-1]
file_name = f"{chat_id}.{msg_id}.{maybe_file_type}"
headers["Content-Disposition"] = f'Content-Disposition: inline; filename="{file_name}"'
if range_header is not None:
start, end = apiutils.get_range_header(range_header, file_size)
size = end - start + 1
headers["content-length"] = str(size)
headers["content-range"] = f"bytes {start}-{end}/{file_size}"
status_code = status.HTTP_206_PARTIAL_CONTENT
return StreamingResponse(
get_msg_media_range_requests(client, msg, start, end),
headers=headers,
status_code=status_code,
)
except Exception as err:
print(f"{err=}")
return Response(f"{err=}", status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
if __name__ == "__main__": if __name__ == "__main__":

72
test.py
View File

@ -4,8 +4,8 @@ import configParse
param = configParse.get_TgToFileSystemParameter() param = configParse.get_TgToFileSystemParameter()
# Remember to use your own values from my.telegram.org! # Remember to use your own values from my.telegram.org!
api_id = param.ApiParameter.api_id api_id = param.tgApi.api_id
api_hash = param.ApiParameter.api_hash api_hash = param.tgApi.api_hash
client = TelegramClient('anon', api_id, api_hash, proxy={ client = TelegramClient('anon', api_id, api_hash, proxy={
'proxy_type': 'socks5', 'proxy_type': 'socks5',
'addr': '172.25.32.1', 'addr': '172.25.32.1',
@ -30,8 +30,12 @@ async def main():
print(me.phone) print(me.phone)
# You can print all the dialogs/conversations that you are part of: # You can print all the dialogs/conversations that you are part of:
async for dialog in client.iter_dialogs(): # async for dialog in client.iter_dialogs():
print(dialog.name, 'has ID', dialog.id) # print(dialog.name, 'has ID', dialog.id)
# test_res = await client.get_input_entity(dialog.id)
# print(test_res)
# await client.send_message(-1001150067822, "test message from python")
# nep_channel = await client.get_dialogs("-1001251458407")
# You can send messages to yourself... # You can send messages to yourself...
# await client.send_message('me', 'Hello, myself!') # await client.send_message('me', 'Hello, myself!')
@ -60,9 +64,14 @@ async def main():
# await client.send_file('me', './test.py') # await client.send_file('me', './test.py')
# You can print the message history of any chat: # You can print the message history of any chat:
message = await client.get_messages('me', ids=206963) # message = await client.get_messages(nep_channel[0])
async for message in client.iter_messages('me'): chat = await client.get_input_entity(-1001216816802)
async for message in client.iter_messages(chat, ids=98724):
print(message.id, message.text) print(message.id, message.text)
# print(message.stringify())
# print(message.to_json())
# print(message.to_dict())
# await client.download_media(message)
# You can download media from messages, too! # You can download media from messages, too!
# The method will return the path where the file was saved. # The method will return the path where the file was saved.
@ -72,3 +81,54 @@ async def main():
with client: with client:
client.loop.run_until_complete(main()) client.loop.run_until_complete(main())
async def start_tg_client(param: configParse.TgToFileSystemParameter):
api_id = param.tgApi.api_id
api_hash = param.tgApi.api_hash
session_name = "test"
proxy_param = {
'proxy_type': param.proxy.proxy_type,
'addr': param.proxy.addr,
'port': param.proxy.port,
} if param.proxy.enable else {}
client = TelegramClient(session_name, api_id, api_hash, proxy=proxy_param)
async def tg_client_main():
# Getting information about yourself
me = await client.get_me()
# "me" is a user object. You can pretty-print
# any Telegram object with the "stringify" method:
print(me.stringify())
# When you print something, you see a representation of it.
# You can access all attributes of Telegram objects with
# the dot operator. For example, to get the username:
username = me.username
print(username)
print(me.phone)
# You can print all the dialogs/conversations that you are part of:
# dialogs = await client.get_dialogs()
# for dialog in dialogs:
# print(f"{dialog.name} has ID {dialog.id}")\
path_task_list = []
async for dialog in client.iter_dialogs():
print(dialog.name, 'has ID', dialog.id)
# path = await client.download_profile_photo(dialog.id)
# t = client.loop.create_task(
# client.download_profile_photo(dialog.id))
# path_task_list.append(t)
# res = await asyncio.gather(*path_task_list)
# for path in res:
# print(path)
# async with client:
# await tg_client_main()
await client.connect()
# qr_login = await client.qr_login()
await client.start()
# print(qr_login.url)
# await qr_login.wait()
await tg_client_main()
await client.disconnect()