103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
# JTFTP - Python/AsyncIO TFTP Server
|
|
# Copyright (C) 2022 Jeffrey C. Ollie
|
|
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
|
|
import io
|
|
import logging
|
|
from typing import Awaitable
|
|
from typing import Callable
|
|
|
|
from jtftp.filesystem import FileMode
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class InMemoryFile:
|
|
filename: bytes
|
|
mode: FileMode
|
|
data: io.BytesIO
|
|
closed: bool
|
|
complete_callback: Callable[[bytes, bytes], Awaitable[None]] | None
|
|
incomplete_callback: Callable[[bytes, bytes], Awaitable[None]] | None
|
|
|
|
def __init__(
|
|
self,
|
|
filename: bytes,
|
|
mode: FileMode,
|
|
*,
|
|
initial_bytes: bytes | None = None,
|
|
complete_callback: Callable[[bytes, bytes], Awaitable[None]],
|
|
incomplete_callback: Callable[[bytes, bytes], Awaitable[None]],
|
|
):
|
|
logger.debug(f"myfile {filename} {mode}")
|
|
self.filename = filename
|
|
self.mode = mode
|
|
self.data = io.BytesIO(initial_bytes)
|
|
self.complete_callback = complete_callback
|
|
self.incomplete_callback = incomplete_callback
|
|
self.closed = False
|
|
|
|
async def length(self) -> int:
|
|
logger.debug("myfile length")
|
|
return len(self.data.getbuffer())
|
|
|
|
async def seek(self, offset: int, whence: int) -> int:
|
|
logger.debug(f"myfile seek {offset} {whence}")
|
|
return self.data.seek(offset, whence)
|
|
|
|
async def read(self, length: int) -> bytes:
|
|
logger.debug(f"myfile read {length}")
|
|
return self.data.read(length)
|
|
|
|
async def write(self, data: bytes) -> int:
|
|
logger.debug(f"myfile write {len(data)}")
|
|
return self.data.write(data)
|
|
|
|
async def close(self, complete: bool) -> None:
|
|
logger.debug(f"myfile close {complete}")
|
|
if complete:
|
|
await self.complete_callback(self.filename, self.data.getvalue())
|
|
else:
|
|
await self.incomplete_callback(self.filename, self.data.getvalue())
|
|
self.closed = True
|
|
|
|
|
|
class InMemoryFilesystem:
|
|
files: dict[bytes, bytes]
|
|
|
|
def __init__(self):
|
|
self.files = {}
|
|
|
|
async def open(self, filename: bytes, mode: FileMode) -> InMemoryFile:
|
|
logger.debug(f"InMemoryFilesystem open {filename} {mode}")
|
|
match mode:
|
|
case FileMode.BINARY_READ:
|
|
if filename not in self.files:
|
|
raise FileNotFoundError
|
|
|
|
return InMemoryFile(filename, mode, initial_bytes=self.files[filename])
|
|
|
|
case FileMode.BINARY_WRITE:
|
|
return InMemoryFile(
|
|
filename, mode, complete_callback=self._complete_callback
|
|
)
|
|
|
|
case _:
|
|
raise ValueError
|
|
|
|
async def _complete_callback(self, filename: bytes, data: bytes):
|
|
self.files[filename] = data
|