Skip to content

Commit

Permalink
fixes logs, restart, reload, upload and some fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
@rking32
rking32 committed Jul 2, 2020
1 parent de58443 commit 514cf5e
Show file tree
Hide file tree
Showing 6 changed files with 178 additions and 145 deletions.
71 changes: 37 additions & 34 deletions userge/core/client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,35 +35,14 @@
_START_TIME = time.time()


class _UsergeBot(Methods, RawClient):
""" UsergeBot, the bot """
def __init__(self, **kwargs) -> None:
_LOG.info(_LOG_STR, "Setting UsergeBot Configs")
super().__init__(session_name=":memory:", **kwargs)
async def _complete_init_tasks() -> None:
if not _INIT_TASKS:
return
await asyncio.gather(*_INIT_TASKS)
_INIT_TASKS.clear()


class Userge(Methods, RawClient):
""" Userge, the userbot """
def __init__(self, **kwargs) -> None:
_LOG.info(_LOG_STR, "Setting Userge Configs")
if not (Config.HU_STRING_SESSION or Config.BOT_TOKEN):
print("Need HU_STRING_SESSION or BOT_TOKEN, Exiting...")
sys.exit()
kwargs = {
'api_id': Config.API_ID,
'api_hash': Config.API_HASH,
'workers': Config.WORKERS
}
if Config.BOT_TOKEN:
if not Config.OWNER_ID:
print("Need OWNER_ID, Exiting...")
sys.exit()
kwargs['bot_token'] = Config.BOT_TOKEN
if Config.HU_STRING_SESSION and Config.BOT_TOKEN:
kwargs['bot'] = _UsergeBot(**kwargs)
kwargs['session_name'] = Config.HU_STRING_SESSION or ":memory:"
super().__init__(**kwargs)

class _AbstractUserge(Methods, RawClient):
@property
def uptime(self) -> str:
""" returns userge uptime """
Expand All @@ -76,15 +55,9 @@ def bot(self) -> '_UsergeBot':
raise UsergeBotNotFound("Need BOT_TOKEN ENV!")
return self._bot

async def _complete_init_tasks(self) -> None:
if not _INIT_TASKS:
return
await asyncio.gather(*_INIT_TASKS)
_INIT_TASKS.clear()

async def finalize_load(self) -> None:
""" finalize the plugins load """
await asyncio.gather(self._complete_init_tasks(), self.manager.init())
await asyncio.gather(_complete_init_tasks(), self.manager.init())

async def load_plugin(self, name: str, reload_plugin: bool = False) -> None:
""" Load plugin to Userge """
Expand Down Expand Up @@ -148,6 +121,36 @@ async def restart(self, update_req: bool = False) -> None: # pylint: disable=ar
os.execl(sys.executable, sys.executable, '-m', 'userge')
sys.exit()


class _UsergeBot(_AbstractUserge):
""" UsergeBot, the bot """
def __init__(self, **kwargs) -> None:
_LOG.info(_LOG_STR, "Setting UsergeBot Configs")
super().__init__(session_name=":memory:", **kwargs)


class Userge(_AbstractUserge):
""" Userge, the userbot """
def __init__(self, **kwargs) -> None:
_LOG.info(_LOG_STR, "Setting Userge Configs")
if not (Config.HU_STRING_SESSION or Config.BOT_TOKEN):
print("Need HU_STRING_SESSION or BOT_TOKEN, Exiting...")
sys.exit()
kwargs = {
'api_id': Config.API_ID,
'api_hash': Config.API_HASH,
'workers': Config.WORKERS
}
if Config.BOT_TOKEN:
if not Config.OWNER_ID:
print("Need OWNER_ID, Exiting...")
sys.exit()
kwargs['bot_token'] = Config.BOT_TOKEN
if Config.HU_STRING_SESSION and Config.BOT_TOKEN:
kwargs['bot'] = _UsergeBot(**kwargs)
kwargs['session_name'] = Config.HU_STRING_SESSION or ":memory:"
super().__init__(**kwargs)

async def start(self) -> None:
""" start client and bot """
_LOG.info(_LOG_STR, "Starting Userge")
Expand Down
14 changes: 6 additions & 8 deletions userge/core/methods/messages/send_as_file.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,7 @@
__all__ = ['SendAsFile']

import os
import asyncio
from typing import Union
from typing import Union, Optional

import aiofiles

Expand All @@ -31,7 +30,7 @@ async def send_as_file(self,
filename: str = "output.txt",
caption: str = '',
log: Union[bool, str] = False,
delete_message: bool = True) -> 'types.bound.Message':
reply_to_message_id: Optional[int] = None) -> 'types.bound.Message':
"""\nYou can send large outputs as file
Example:
Expand Down Expand Up @@ -63,24 +62,23 @@ async def send_as_file(self,
If ``True``, the message will be deleted
after sending the file.
reply_to_message_id (``int``, *optional*):
If the message is a reply, ID of the original message.
Returns:
On success, the sent Message is returned.
"""
async with aiofiles.open(filename, "w+", encoding="utf8") as out_file:
await out_file.write(text)
reply_to_id = self.reply_to_message.message_id if self.reply_to_message \
else self.message_id
_LOG.debug(_LOG_STR, f"Uploading {filename} To Telegram")
msg = await self.send_document(chat_id=chat_id,
document=filename,
caption=caption,
disable_notification=True,
reply_to_message_id=reply_to_id)
reply_to_message_id=reply_to_message_id)
os.remove(filename)
if log:
if isinstance(log, str):
self._channel.update(log)
await self._channel.fwd_msg(msg)
if delete_message:
asyncio.get_event_loop().create_task(self.delete())
return types.bound.Message(self, msg)
27 changes: 19 additions & 8 deletions userge/core/types/bound/message.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
__all__ = ['Message']

import re
import asyncio
from typing import List, Dict, Union, Optional, Sequence

from pyrogram import InlineKeyboardMarkup, Message as RawMessage
Expand All @@ -32,7 +33,7 @@ def _msg_to_dict(message: RawMessage) -> Dict[str, object]:
kwargs_ = vars(message)
del message
for key_ in ['_client', '_channel', '_filtered', '_process_canceled',
'_filtered_input_str', '_flags', '_kwargs']:
'client', '_filtered_input_str', '_flags', '_kwargs']:
if key_ in kwargs_:
del kwargs_[key_]
return kwargs_
Expand All @@ -41,7 +42,7 @@ def _msg_to_dict(message: RawMessage) -> Dict[str, object]:
class Message(RawMessage):
""" Modded Message Class For Userge """
def __init__(self,
client: '_client.Userge',
client: Union['_client.Userge', '_client.UsergeBot'],
message: RawMessage,
**kwargs: Union[str, bool]) -> None:
super().__init__(client=client, **_msg_to_dict(message))
Expand All @@ -56,6 +57,11 @@ def __init__(self,
self._flags: Dict[str, str] = {}
self._kwargs = kwargs

@property
def client(self) -> Union['_client.Userge', '_client.UsergeBot']:
""" returns client """
return self._client

@property
def input_str(self) -> str:
""" Returns the input string without command """
Expand Down Expand Up @@ -148,12 +154,17 @@ async def send_as_file(self,
Returns:
On success, the sent Message is returned.
"""
return self._client.send_as_file(chat_id=self.chat.id,
text=text,
filename=filename,
caption=caption,
log=log,
delete_message=delete_message)
reply_to_id = self.reply_to_message.message_id if self.reply_to_message \
else self.message_id
if delete_message:
asyncio.get_event_loop().create_task(self.delete())
return await self._client.send_as_file(chat_id=self.chat.id,
text=text,
filename=filename,
caption=caption,
log=log,
delete_message=delete_message,
reply_to_message_id=reply_to_id)

async def reply(self,
text: str,
Expand Down
127 changes: 73 additions & 54 deletions userge/core/types/new/channel_logger.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from typing import Optional, Tuple

from pyrogram import Message as RawMessage
from pyrogram.errors.exceptions.bad_request_400 import ChannelInvalid

from userge import logging, Config
from userge.utils import SafeDict
Expand Down Expand Up @@ -92,9 +93,12 @@ async def log(self, text: str) -> Optional[int]:
"""
_LOG.debug(_LOG_STR, f"logging text : {text} to channel : {Config.LOG_CHANNEL_ID}")
if Config.LOG_CHANNEL_ID:
msg = await self._client.send_message(chat_id=Config.LOG_CHANNEL_ID,
text=self._string.format(text.strip()))
return msg.message_id
try:
msg = await self._client.send_message(chat_id=Config.LOG_CHANNEL_ID,
text=self._string.format(text.strip()))
return msg.message_id
except ChannelInvalid:
pass

async def fwd_msg(self,
message: '_message.Message',
Expand Down Expand Up @@ -124,15 +128,19 @@ async def fwd_msg(self,
_LOG.debug(
_LOG_STR, f"forwarding msg : {message} to channel : {Config.LOG_CHANNEL_ID}")
if Config.LOG_CHANNEL_ID and isinstance(message, RawMessage):
if message.media:
asyncio.get_event_loop().create_task(self.log("**Forwarding Message...**"))
await self._client.forward_messages(chat_id=Config.LOG_CHANNEL_ID,
from_chat_id=message.chat.id,
message_ids=message.message_id,
as_copy=as_copy,
remove_caption=remove_caption)
else:
await self.log(message.text.html)
try:
if message.media:
asyncio.get_event_loop().create_task(self.log("**Forwarding Message...**"))
await self._client.forward_messages(chat_id=Config.LOG_CHANNEL_ID,
from_chat_id=message.chat.id,
message_ids=message.message_id,
as_copy=as_copy,
remove_caption=remove_caption)
else:
await self.log(
message.text.html if hasattr(message.text, 'html') else message.text)
except ChannelInvalid:
pass

async def store(self,
message: Optional['_message.Message'],
Expand All @@ -149,21 +157,25 @@ async def store(self,
Returns:
message_id on success or None
"""
caption = caption or ''
if message and message.caption:
caption = caption + message.caption.html
if message and message.media:
if caption:
caption = self._string.format(caption.strip())
file_id, file_ref = _get_file_id_and_ref(message)
msg = await self._client.send_cached_media(chat_id=Config.LOG_CHANNEL_ID,
file_id=file_id,
file_ref=file_ref,
caption=caption)
message_id = msg.message_id
else:
message_id = await self.log(caption)
return message_id
if Config.LOG_CHANNEL_ID:
caption = caption or ''
if message and message.caption:
caption = caption + message.caption.html
if message and message.media:
if caption:
caption = self._string.format(caption.strip())
file_id, file_ref = _get_file_id_and_ref(message)
try:
msg = await self._client.send_cached_media(chat_id=Config.LOG_CHANNEL_ID,
file_id=file_id,
file_ref=file_ref,
caption=caption)
message_id = msg.message_id
except ChannelInvalid:
message_id = None
else:
message_id = await self.log(caption)
return message_id

async def forward_stored(self,
message_id: int,
Expand Down Expand Up @@ -192,30 +204,37 @@ async def forward_stored(self,
Returns:
None
"""
message = await self._client.get_messages(chat_id=Config.LOG_CHANNEL_ID,
message_ids=message_id)
caption = ''
if message.caption:
caption = message.caption.html.split('\n\n', maxsplit=1)[-1]
elif message.text:
caption = message.text.html.split('\n\n', maxsplit=1)[-1]
if caption:
u_dict = await self._client.get_user_dict(user_id)
chat = await self._client.get_chat(chat_id)
u_dict.update(
{'chat': chat.title if chat.title else "this group", 'count': chat.members_count})
caption = caption.format_map(SafeDict(**u_dict))
if message.media:
file_id, file_ref = _get_file_id_and_ref(message)
msg = await self._client.send_cached_media(chat_id=chat_id,
file_id=file_id,
file_ref=file_ref,
caption=caption,
reply_to_message_id=reply_to_message_id)
else:
msg = await self._client.send_message(chat_id=chat_id,
text=caption,
reply_to_message_id=reply_to_message_id)
if del_in and msg:
await asyncio.sleep(del_in)
await msg.delete()
if Config.LOG_CHANNEL_ID:
try:
message = await self._client.get_messages(chat_id=Config.LOG_CHANNEL_ID,
message_ids=message_id)
caption = ''
if message.caption:
caption = message.caption.html.split('\n\n', maxsplit=1)[-1]
elif message.text:
caption = message.text.html.split('\n\n', maxsplit=1)[-1]
if caption:
u_dict = await self._client.get_user_dict(user_id)
chat = await self._client.get_chat(chat_id)
u_dict.update(
{'chat': chat.title if chat.title else "this group",
'count': chat.members_count})
caption = caption.format_map(SafeDict(**u_dict))
if message.media:
file_id, file_ref = _get_file_id_and_ref(message)
msg = await self._client.send_cached_media(
chat_id=chat_id,
file_id=file_id,
file_ref=file_ref,
caption=caption,
reply_to_message_id=reply_to_message_id)
else:
msg = await self._client.send_message(
chat_id=chat_id,
text=caption,
reply_to_message_id=reply_to_message_id)
if del_in and msg:
await asyncio.sleep(del_in)
await msg.delete()
except ChannelInvalid:
pass
Loading

0 comments on commit 514cf5e

Please sign in to comment.