nikita
16b849ebd8
- Added example `file_bot.py` - Added *.tmp to .gitignore - Added CTX_FILE_SIZE, CTX_FILE_MIME and CTX_FILE_NAME - Added MessageHasFile filter - Added `send_file` and `send_file_bytes` for ClientSender and Bot
97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
"""
|
|
This example implements Matrix bot that calculates SHA256 for a file sent by
|
|
user.
|
|
|
|
It uses environment variables to specify authorization data. Use Ctrl+C to stop
|
|
the bot.
|
|
"""
|
|
|
|
import asyncio
|
|
import aiofiles
|
|
import os
|
|
import logging
|
|
import hashlib
|
|
|
|
from mab import *
|
|
|
|
from _environment import check_environment
|
|
|
|
async def on_text_message(ctx: EventContext) -> None:
|
|
"""This callback is called when a text message arrives."""
|
|
await ctx.bot.send_text(ctx.room, "Please send a file/image/video")
|
|
|
|
async def on_file_message(ctx: EventContext) -> None:
|
|
"""This callback is called when a file arrives."""
|
|
sha = hashlib.sha256()
|
|
does_temp_exist = False
|
|
if ctx[CTX_FILE_SIZE] > 1_000_000:
|
|
await ctx.bot.send_text(
|
|
ctx.room,
|
|
"The file is larger than 1 MB, downloading to filesystem"
|
|
)
|
|
await ctx.bot.download_file(ctx, path="temp.tmp")
|
|
async with aiofiles.open("temp.tmp", "rb") as f:
|
|
while True:
|
|
data = await f.read(64 * 1024)
|
|
if not data:
|
|
break
|
|
sha.update(data)
|
|
does_temp_exist = True
|
|
else:
|
|
await ctx.bot.send_text(
|
|
ctx.room,
|
|
"The file is smaller than 1 MB, downloading to RAM"
|
|
)
|
|
content = await ctx.bot.download_file(ctx, path=None)
|
|
sha.update(content)
|
|
# result
|
|
response = f"SHA256 for file `{ctx[CTX_FILE_NAME]}`"
|
|
response += f" ({ctx[CTX_FILE_SIZE]} bytes, {ctx[CTX_FILE_MIME]})"
|
|
await ctx.bot.send_file_bytes(
|
|
room=ctx.room,
|
|
data=sha.hexdigest().encode("utf-8"),
|
|
filename="hash of the file.txt",
|
|
mime_type="text/plain",
|
|
text=response
|
|
)
|
|
# resend the file to test uploading
|
|
if does_temp_exist:
|
|
await ctx.bot.send_file(
|
|
room=ctx.room,
|
|
path="temp.tmp",
|
|
filename=ctx[CTX_FILE_NAME],
|
|
text="This is the file you have sent, but it was reuploaded"
|
|
)
|
|
try:
|
|
os.unlink("temp.tmp")
|
|
except:
|
|
pass
|
|
|
|
|
|
async def main() -> None:
|
|
"""Application entry point"""
|
|
logging.basicConfig(level=logging.INFO)
|
|
logging.getLogger("nio").setLevel(logging.CRITICAL+1)
|
|
check_environment()
|
|
|
|
config = MatrixBotConfig(
|
|
matrix_homeserver_url=os.environ["MATRIX_HOMESERVER"],
|
|
matrix_username_localpart=os.environ["MATRIX_USERNAME"],
|
|
storage_directory="session_storage"
|
|
)
|
|
bot = MatrixBot(config)
|
|
bot.add_callback(
|
|
~SenderIsBotFilter() & BodyExistsFilter() & MessageTypeFilter(MessageType.TEXT),
|
|
on_text_message)
|
|
bot.add_callback(
|
|
~SenderIsBotFilter() & MessageHasFile() & NewMessageFilter(),
|
|
on_file_message)
|
|
|
|
# run until Ctrl+C
|
|
try:
|
|
await bot.run()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main()) |