Roman Donchenko
4aa0be3df5
Currently, all source files are located in `/home/django`, and owned by `django`. This means that if there's any vulnerability that lets an attacker overwrite files in the server, they can replace source files with their own code, and potentially get that code executed. That's pretty bad, so I want to harden against that. Make all source files owned by root, and move them to `/opt/cvat`. Add a `manage.py` symlink in `/home/django` for backwards compatibility. It happens that if a script is a symlink, Python does not add the symlink's directory to `sys.path`, which is great for us, since that lets us avoid a writable directory on there. Still, even though `/home/django/manage.py` is owned by root, an attacker could potentially be able to delete it and replace it with their own malicious file. To be a bit more safe, replace `~/manage.py` calls in backend scripts with `django-admin`. To make sure CVAT can still find the data directory, add a new environment variable, `CVAT_BASE_DIR` and set it in the Docker image. This also fixes a minor bug: we no longer override the `HOME` environment variable in the `Dockerfile`, so now it's automatically set by `Docker` depending on the current user.
132 lines
3.3 KiB
Python
132 lines
3.3 KiB
Python
# Copyright (C) CVAT.ai Corporation
|
|
#
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
import io
|
|
import subprocess
|
|
from collections.abc import Generator
|
|
from fractions import Fraction
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
import av
|
|
import av.video.reformatter
|
|
import librosa
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from shared.fixtures.init import get_server_image_tag
|
|
|
|
|
|
def generate_image_file(filename="image.png", size=(100, 50), color=(0, 0, 0)):
|
|
f = BytesIO()
|
|
f.name = filename
|
|
image = Image.new("RGB", size=size, color=color)
|
|
image.save(f)
|
|
f.seek(0)
|
|
|
|
return f
|
|
|
|
|
|
def generate_image_files(
|
|
count: int,
|
|
*,
|
|
prefixes: list[str] | None = None,
|
|
filenames: list[str] | None = None,
|
|
sizes: list[tuple[int, int]] | None = None,
|
|
) -> list[BytesIO]:
|
|
assert not (prefixes and filenames), "prefixes cannot be used together with filenames"
|
|
assert not prefixes or len(prefixes) == count
|
|
assert not filenames or len(filenames) == count
|
|
|
|
images = []
|
|
for i in range(count):
|
|
prefix = prefixes[i] if prefixes else ""
|
|
filename = f"{prefix}{i}.jpeg" if not filenames else filenames[i]
|
|
image = generate_image_file(
|
|
filename, color=(i, i, i), **({"size": sizes[i]}) if sizes else {}
|
|
)
|
|
images.append(image)
|
|
|
|
return images
|
|
|
|
|
|
def generate_video_file(
|
|
num_frames: int,
|
|
*,
|
|
size: tuple[int, int] = (100, 50),
|
|
invalid_keyframes: bool = False,
|
|
) -> BytesIO:
|
|
f = BytesIO()
|
|
f.name = "video.mkv"
|
|
chapters = [
|
|
{
|
|
"id": 0,
|
|
"start": 0,
|
|
"end": 100,
|
|
"time_base": Fraction(1, 1000),
|
|
"metadata": {"title": "Intro"},
|
|
}
|
|
]
|
|
|
|
with av.open(f, "w") as container:
|
|
container.set_chapters(chapters)
|
|
stream = container.add_stream("mjpeg", rate=60)
|
|
stream.width = size[0]
|
|
stream.height = size[1]
|
|
stream.color_range = av.video.reformatter.ColorRange.JPEG
|
|
|
|
for i in range(num_frames):
|
|
frame = av.VideoFrame.from_image(Image.new("RGB", size=size, color=(i, i, i)))
|
|
for packet in stream.encode(frame):
|
|
if invalid_keyframes:
|
|
# Specify pts/dts values that result in 0 valid keyframes
|
|
packet.pts = 10
|
|
packet.dts = 10
|
|
|
|
container.mux(packet)
|
|
|
|
f.seek(0)
|
|
|
|
return f
|
|
|
|
|
|
def read_video_file(file: BytesIO) -> Generator[Image.Image, None, None]:
|
|
file.seek(0)
|
|
|
|
with av.open(file) as container:
|
|
video_stream = container.streams.video[0]
|
|
|
|
for frame in container.decode(video_stream):
|
|
yield frame.to_image()
|
|
|
|
|
|
def generate_manifest(path: str) -> None:
|
|
command = [
|
|
"docker",
|
|
"run",
|
|
"--rm",
|
|
"-u",
|
|
"root:root",
|
|
"-v",
|
|
f"{path}:/local",
|
|
"--entrypoint",
|
|
"python3",
|
|
get_server_image_tag(),
|
|
"/opt/cvat/utils/dataset_manifest/create.py",
|
|
"--output-dir",
|
|
"/local",
|
|
"/local",
|
|
]
|
|
try:
|
|
subprocess.check_output(command, stderr=subprocess.PIPE)
|
|
except subprocess.CalledProcessError as e:
|
|
print(e.stderr.decode("utf-8"))
|
|
raise
|
|
|
|
|
|
def read_audio_pcm(
|
|
f: Path | io.IOBase, *, offset_ms: int = 0, rate: int = 8000
|
|
) -> tuple[np.ndarray, float]:
|
|
return librosa.load(f, mono=True, sr=rate, offset=offset_ms / 1000)
|