Skip to content

Commit

Permalink
feat: add file upload support and utility functions to client library
Browse files Browse the repository at this point in the history
This update adds file uploading support to the Dioptra client library through a new DioptraFile
class, supporting utility functions, and extending the client's post method to accept `data` and
`files` arguments. The changes enable users to upload single files, multiple files, or all files
within a directory while enforcing strict path validation in the associated filenames to prevent
server-side directory traversal attacks.

The new functionality includes:

- A DioptraFile class that validates filenames to ensure they are normalized, POSIX-compliant
  relative paths without traversal components (e.g., "..").
- File selection utility functions:
  - select_files_in_directory: Recursively selects files from a directory with optional regex
    pattern filtering
  - select_one_or_more_files: Selects individual files with optional renaming support to handle
    filename collisions
- Extended POST request in the client to handle multipart form data that contains regular form
  fields and file uploads. All file uploads, whether they're single file or multiple files, are
  streamed to the server.

The DioptraSession implementation for the requests-based client and the Flask test client have both
been updated to support the new functionality.
  • Loading branch information
jkglasbrenner committed Jan 15, 2025
1 parent 0140f38 commit db9b5a5
Show file tree
Hide file tree
Showing 6 changed files with 682 additions and 6 deletions.
5 changes: 5 additions & 0 deletions src/dioptra/client/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,14 +14,19 @@
#
# ACCESS THE FULL CC BY 4.0 LICENSE HERE:
# https://creativecommons.org/licenses/by/4.0/legalcode
from .base import DioptraFile
from .client import (
DioptraClient,
connect_json_dioptra_client,
connect_response_dioptra_client,
)
from .utils import select_files_in_directory, select_one_or_more_files

__all__ = [
"connect_response_dioptra_client",
"connect_json_dioptra_client",
"select_files_in_directory",
"select_one_or_more_files",
"DioptraClient",
"DioptraFile",
]
79 changes: 77 additions & 2 deletions src/dioptra/client/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,19 @@
#
# ACCESS THE FULL CC BY 4.0 LICENSE HERE:
# https://creativecommons.org/licenses/by/4.0/legalcode
import posixpath
import re
from abc import ABC, abstractmethod
from pathlib import Path
from dataclasses import dataclass
from io import BufferedReader
from pathlib import Path, PurePosixPath, PureWindowsPath
from posixpath import join as urljoin
from typing import Any, ClassVar, Generic, Protocol, TypeVar

T = TypeVar("T")

DOTS_REGEX = re.compile(r"^\.\.\.+$")


class DioptraClientError(Exception):
"""Base class for client errors"""
Expand Down Expand Up @@ -91,6 +97,52 @@ def json(self) -> dict[str, Any]:
... # fmt: skip


@dataclass
class DioptraFile(object):
"""A file to be uploaded to the Dioptra API.
Attributes:
filename: The name of the file.
stream: The file stream.
content_type: The content type of the file.
"""

filename: str
stream: BufferedReader
content_type: str | None

def __post_init__(self) -> None:
if PureWindowsPath(self.filename).as_posix() != str(PurePosixPath(self.filename)): # noqa: B950; fmt: skip
raise ValueError(
"Invalid filename (reason: filename is a Windows path): "
f"{self.filename}"
)

if posixpath.normpath(self.filename) != self.filename:
raise ValueError(
"Invalid filename (reason: filename is not normalized): "
f"{self.filename}"
)

if not PurePosixPath(self.filename).is_relative_to("."):
raise ValueError(
"Invalid filename (reason: filename is not relative to ./): "
f"{self.filename}"
)

if PurePosixPath("..") in PurePosixPath(posixpath.normpath(self.filename)).parents: # noqa: B950; fmt: skip
raise ValueError(
"Invalid filename (reason: filename is not a sub-directory of ./): "
f"{self.filename}"
)

if any([DOTS_REGEX.match(str(x)) for x in PurePosixPath(posixpath.normpath(self.filename)).parts]): # noqa: B950; fmt: skip
raise ValueError(
"Invalid filename (reason: filename contains a sub-directory name that "
f"is all dots): {self.filename}"
)


class DioptraSession(ABC, Generic[T]):
"""The interface for communicating with the Dioptra API."""

Expand All @@ -117,6 +169,8 @@ def make_request(
url: str,
params: dict[str, Any] | None = None,
json_: dict[str, Any] | None = None,
data: dict[str, Any] | None = None,
files: dict[str, DioptraFile | list[DioptraFile]] | None = None,
) -> DioptraResponseProtocol:
"""Make a request to the API.
Expand All @@ -129,6 +183,10 @@ def make_request(
params: The query parameters to include in the request. Optional, defaults
to None.
json_: The JSON data to include in the request. Optional, defaults to None.
data: A dictionary to send in the body of the request as part of a
multipart form. Optional, defaults to None.
files: Dictionary of "name": DioptraFile or lists of DioptraFile pairs to be
uploaded. Optional, defaults to None.
Returns:
The response from the API.
Expand Down Expand Up @@ -179,6 +237,8 @@ def post(
*parts,
params: dict[str, Any] | None = None,
json_: dict[str, Any] | None = None,
data: dict[str, Any] | None = None,
files: dict[str, DioptraFile | list[DioptraFile]] | None = None,
) -> T:
"""Make a POST request to the API.
Expand All @@ -188,6 +248,10 @@ def post(
params: The query parameters to include in the request. Optional, defaults
to None.
json_: The JSON data to include in the request. Optional, defaults to None.
data: A dictionary to send in the body of the request as part of a
multipart form. Optional, defaults to None.
files: Dictionary of "name": DioptraFile or lists of DioptraFile pairs to be
uploaded. Optional, defaults to None.
Returns:
The response from the API.
Expand Down Expand Up @@ -311,6 +375,8 @@ def _post(
*parts,
params: dict[str, Any] | None = None,
json_: dict[str, Any] | None = None,
data: dict[str, Any] | None = None,
files: dict[str, DioptraFile | list[DioptraFile]] | None = None,
) -> DioptraResponseProtocol:
"""Make a POST request to the API.
Expand All @@ -323,12 +389,21 @@ def _post(
params: The query parameters to include in the request. Optional, defaults
to None.
json_: The JSON data to include in the request. Optional, defaults to None.
data: A dictionary to send in the body of the request as part of a
multipart form. Optional, defaults to None.
files: Dictionary of "name": DioptraFile or lists of DioptraFile pairs to be
uploaded. Optional, defaults to None.
Returns:
A response object that implements the DioptraResponseProtocol interface.
"""
return self.make_request(
"post", self.build_url(endpoint, *parts), params=params, json_=json_
"post",
self.build_url(endpoint, *parts),
params=params,
json_=json_,
data=data,
files=files,
)

def _delete(
Expand Down
Loading

0 comments on commit db9b5a5

Please sign in to comment.