Added file uploading support, v0.5.2
- 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
This commit is contained in:
3
.gitignore
vendored
3
.gitignore
vendored
@@ -5,4 +5,5 @@ session_storage/
|
|||||||
dist/
|
dist/
|
||||||
*.egg-info/
|
*.egg-info/
|
||||||
*.swp
|
*.swp
|
||||||
*.swo
|
*.swo
|
||||||
|
*.tmp
|
||||||
97
examples/file_bot.py
Normal file
97
examples/file_bot.py
Normal file
@@ -0,0 +1,97 @@
|
|||||||
|
"""
|
||||||
|
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())
|
||||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|||||||
|
|
||||||
[project]
|
[project]
|
||||||
name = "mab"
|
name = "mab"
|
||||||
version = "0.5.1"
|
version = "0.5.2"
|
||||||
authors = [
|
authors = [
|
||||||
{ name = "Tyukalov Nikita", email = "nikita@tyukalov.su" }
|
{ name = "Tyukalov Nikita", email = "nikita@tyukalov.su" }
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -27,6 +27,9 @@ __all__ = [
|
|||||||
"CTX_CMD_PREFIX",
|
"CTX_CMD_PREFIX",
|
||||||
"CTX_CMD_VERB",
|
"CTX_CMD_VERB",
|
||||||
"CTX_CMD_ARGS",
|
"CTX_CMD_ARGS",
|
||||||
|
"CTX_FILE_SIZE",
|
||||||
|
"CTX_FILE_MIME",
|
||||||
|
"CTX_FILE_NAME",
|
||||||
|
|
||||||
# .bot
|
# .bot
|
||||||
"MatrixBot",
|
"MatrixBot",
|
||||||
@@ -50,4 +53,5 @@ __all__ = [
|
|||||||
"RedactedMessageFilter",
|
"RedactedMessageFilter",
|
||||||
"SenderIsFilter",
|
"SenderIsFilter",
|
||||||
"SenderIsBotFilter",
|
"SenderIsBotFilter",
|
||||||
|
"MessageHasFile",
|
||||||
]
|
]
|
||||||
@@ -361,4 +361,123 @@ class ClientSender:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
# send
|
# send
|
||||||
|
return (await self.send_content(room, content)).event_id
|
||||||
|
|
||||||
|
async def send_file(self,
|
||||||
|
room: MatrixRoom | str,
|
||||||
|
path: Path | str,
|
||||||
|
*,
|
||||||
|
filename: str | None = None,
|
||||||
|
mime_type: str | None = None,
|
||||||
|
text: str | None = None,
|
||||||
|
is_html: bool | None = None,
|
||||||
|
timeout: float | None = 60 * 60):
|
||||||
|
"""
|
||||||
|
Send the file to `room`. Please note that formatted text is displayed
|
||||||
|
incorrectly in some clients as of September 8th, 2026.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- room - the room to send the file to
|
||||||
|
- path - path to the file
|
||||||
|
- filename - filename to use for upload (`None` to use basename from
|
||||||
|
`path`)
|
||||||
|
- mime_type - mime type to use (`None` for autodetect using content
|
||||||
|
of `path`)
|
||||||
|
- text - caption to use (`None` to disable)
|
||||||
|
- is_html - whether the text is HTML-formatted (`None` for auto)
|
||||||
|
- timeout - upload timeout in seconds (`None` to disable)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- `event_id` of sent message on success
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
if self._client is None or self._config is None:
|
||||||
|
raise RuntimeError("ClientSender is not set up")
|
||||||
|
# get filename if not set
|
||||||
|
if not filename:
|
||||||
|
filename = os.path.basename(path)
|
||||||
|
# caption must not actually be empty
|
||||||
|
if text is None or not text.strip():
|
||||||
|
text = filename
|
||||||
|
is_html = False
|
||||||
|
# get the mime type if not specified
|
||||||
|
if mime_type is None:
|
||||||
|
mime_type = magic.from_file(path, mime=True)
|
||||||
|
# upload
|
||||||
|
async with asyncio.timeout(timeout):
|
||||||
|
upload_result = await self._uploader.upload_file(
|
||||||
|
path, mime_type=mime_type, filename=filename)
|
||||||
|
# send
|
||||||
|
content = {
|
||||||
|
"msgtype": "m.file",
|
||||||
|
"filename": filename,
|
||||||
|
**self._process_html_text(text, is_html),
|
||||||
|
"file": {
|
||||||
|
"url": upload_result.response.content_uri,
|
||||||
|
"mimetype": mime_type,
|
||||||
|
**upload_result.keys
|
||||||
|
},
|
||||||
|
"info": {
|
||||||
|
"mimetype": mime_type,
|
||||||
|
"size": upload_result.filesize
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return (await self.send_content(room, content)).event_id
|
||||||
|
|
||||||
|
async def send_file_bytes(self,
|
||||||
|
room: MatrixRoom | str,
|
||||||
|
data: bytes,
|
||||||
|
filename: str,
|
||||||
|
*,
|
||||||
|
mime_type: str | None = None,
|
||||||
|
text: str | None = None,
|
||||||
|
is_html: bool | None = None,
|
||||||
|
timeout: float | None = 60 * 60) -> str:
|
||||||
|
"""
|
||||||
|
Send the file to `room`.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- room - the room to send the file to
|
||||||
|
- data - the content of the file to send
|
||||||
|
- filename - filename to use for the file
|
||||||
|
- mime_type - mime type to use (`None` for autodetect using `data`
|
||||||
|
content)
|
||||||
|
- text - caption to use (`None` to disable)
|
||||||
|
- is_html - whether the text is HTML-formatted (`None` for auto)
|
||||||
|
- timeout - upload timeout in seconds (`None` to disable)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- `event_id` of sent message on success
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
# caption must not actually be empty
|
||||||
|
if text is None or not text.strip():
|
||||||
|
text = filename
|
||||||
|
is_html = False
|
||||||
|
# check if the file is image
|
||||||
|
if not mime_type:
|
||||||
|
mime_type = magic.from_buffer(data, mime=True)
|
||||||
|
# upload
|
||||||
|
buffer = BytesIO(data)
|
||||||
|
async with asyncio.timeout(timeout):
|
||||||
|
upload_result = await self._uploader.upload_using_provider(
|
||||||
|
provider=buffer,
|
||||||
|
mime_type=mime_type,
|
||||||
|
filename=filename,
|
||||||
|
filesize=len(data))
|
||||||
|
# prepare the content and send
|
||||||
|
content = {
|
||||||
|
"msgtype": "m.file",
|
||||||
|
"filename": filename,
|
||||||
|
**self._process_html_text(text, is_html),
|
||||||
|
"file": {
|
||||||
|
"url": upload_result.response.content_uri,
|
||||||
|
"mimetype": mime_type,
|
||||||
|
**upload_result.keys
|
||||||
|
},
|
||||||
|
"info": {
|
||||||
|
"mimetype": mime_type,
|
||||||
|
"size": upload_result.filesize
|
||||||
|
}
|
||||||
|
}
|
||||||
return (await self.send_content(room, content)).event_id
|
return (await self.send_content(room, content)).event_id
|
||||||
@@ -252,6 +252,77 @@ class MatrixBot:
|
|||||||
timeout=timeout
|
timeout=timeout
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def send_file(self,
|
||||||
|
room: MatrixRoom | str,
|
||||||
|
path: Path | str,
|
||||||
|
*,
|
||||||
|
filename: str | None = None,
|
||||||
|
mime_type: str | None = None,
|
||||||
|
text: str | None = None,
|
||||||
|
is_html: bool | None = None,
|
||||||
|
timeout: float | None = 60 * 60):
|
||||||
|
"""
|
||||||
|
Send the file to `room`. Please note that formatted text is displayed
|
||||||
|
incorrectly in some clients as of September 8th, 2026.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- room - the room to send the file to
|
||||||
|
- path - path to the file
|
||||||
|
- filename - filename to use for upload (`None` to use basename from
|
||||||
|
`path`)
|
||||||
|
- mime_type - mime type to use (`None` for auto)
|
||||||
|
- text - caption to use (`None` to disable)
|
||||||
|
- is_html - whether the text is HTML-formatted (`None` for auto)
|
||||||
|
- timeout - upload timeout in seconds (`None` to disable)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- `event_id` of sent message on success
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
return await self._client_sender.send_file(
|
||||||
|
room=room,
|
||||||
|
path=path,
|
||||||
|
filename=filename,
|
||||||
|
mime_type=mime_type,
|
||||||
|
text=text,
|
||||||
|
is_html=is_html,
|
||||||
|
timeout=timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
async def send_file_bytes(self,
|
||||||
|
room: MatrixRoom | str,
|
||||||
|
data: bytes,
|
||||||
|
filename: str,
|
||||||
|
*,
|
||||||
|
mime_type: str | None = None,
|
||||||
|
text: str | None = None,
|
||||||
|
is_html: bool | None = None,
|
||||||
|
timeout: float | None = 60 * 60) -> str:
|
||||||
|
"""
|
||||||
|
Send the file to `room`.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- room - the room to send the file to
|
||||||
|
- data - the content of the file to send
|
||||||
|
- filename - filename to use for the file
|
||||||
|
- text - caption to use (`None` to disable)
|
||||||
|
- is_html - whether the text is HTML-formatted (`None` for auto)
|
||||||
|
- timeout - upload timeout in seconds (`None` to disable)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- `event_id` of sent message on success
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
return await self._client_sender.send_file_bytes(
|
||||||
|
room=room,
|
||||||
|
data=data,
|
||||||
|
filename=filename,
|
||||||
|
mime_type=mime_type,
|
||||||
|
text=text,
|
||||||
|
is_html=is_html,
|
||||||
|
timeout=timeout
|
||||||
|
)
|
||||||
|
|
||||||
@overload
|
@overload
|
||||||
async def download_file(self,
|
async def download_file(self,
|
||||||
source: EventContext | Event,
|
source: EventContext | Event,
|
||||||
|
|||||||
@@ -34,6 +34,15 @@ CTX_CMD_ARGS = ContextDataKey[list[str]]("CTX_CMD_ARGS")
|
|||||||
CTX_ROOM_ENCRYPTED = ContextDataKey[bool]("CTX_ROOM_ENCRYPTED")
|
CTX_ROOM_ENCRYPTED = ContextDataKey[bool]("CTX_ROOM_ENCRYPTED")
|
||||||
"""True if the room is encrypted"""
|
"""True if the room is encrypted"""
|
||||||
|
|
||||||
|
CTX_FILE_SIZE = ContextDataKey[int]("CTX_FILE_SIZE")
|
||||||
|
"""Size of the file attached to the message (bytes)"""
|
||||||
|
|
||||||
|
CTX_FILE_MIME = ContextDataKey[str]("CTX_FILE_MIME")
|
||||||
|
"""Mime type of the file attached to the message"""
|
||||||
|
|
||||||
|
CTX_FILE_NAME = ContextDataKey[str]("CTX_FILE_NAME")
|
||||||
|
"""Name of the file attached to the message"""
|
||||||
|
|
||||||
|
|
||||||
#
|
#
|
||||||
# EventContext implementation
|
# EventContext implementation
|
||||||
|
|||||||
@@ -1,7 +1,12 @@
|
|||||||
import traceback
|
import traceback
|
||||||
from .base import BaseEventFilter, EventTypeFilter
|
from .base import BaseEventFilter, EventTypeFilter
|
||||||
from ..types import MessageType
|
from ..types import MessageType
|
||||||
from ..context import EventContext, CTX_MESSAGE_TYPE, CTX_SENDER
|
from ..context import (EventContext,
|
||||||
|
CTX_MESSAGE_TYPE,
|
||||||
|
CTX_SENDER,
|
||||||
|
CTX_FILE_SIZE,
|
||||||
|
CTX_FILE_MIME,
|
||||||
|
CTX_FILE_NAME)
|
||||||
|
|
||||||
from nio import AsyncClient
|
from nio import AsyncClient
|
||||||
from nio import MatrixRoom, Event
|
from nio import MatrixRoom, Event
|
||||||
@@ -118,4 +123,29 @@ class SenderIsBotFilter(BaseEventFilter):
|
|||||||
async def __call__(self, context: EventContext) -> bool:
|
async def __call__(self, context: EventContext) -> bool:
|
||||||
if not await super().__call__(context):
|
if not await super().__call__(context):
|
||||||
return False
|
return False
|
||||||
return context.bot.get_client().user_id == context.event.sender
|
return context.bot.get_client().user_id == context.event.sender
|
||||||
|
|
||||||
|
class MessageHasFile(BaseEventFilter):
|
||||||
|
"""
|
||||||
|
This filter returns True if the message contains file that can be
|
||||||
|
downloaded.
|
||||||
|
|
||||||
|
This filter sets `CTX_FILE_NAME`, `CTX_FILE_SIZE` and `CTX_FILE_MIME`
|
||||||
|
variables in the context (if they are present in the).
|
||||||
|
"""
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
|
||||||
|
async def __call__(self, context: EventContext) -> bool:
|
||||||
|
if not await super().__call__(context):
|
||||||
|
return False
|
||||||
|
content: dict | None = context.event.source.get("content")
|
||||||
|
if content is None:
|
||||||
|
return False
|
||||||
|
# `file` if encrypted, `url` if not encrypted
|
||||||
|
if "file" not in content and "url" not in content:
|
||||||
|
return False
|
||||||
|
context[CTX_FILE_NAME] = content.get("filename") or content.get("body")
|
||||||
|
context[CTX_FILE_SIZE] = content["info"].get("size")
|
||||||
|
context[CTX_FILE_MIME] = content["info"].get("mimetype")
|
||||||
|
return True
|
||||||
Reference in New Issue
Block a user