JOJ3-config-generator/joj3_config_generator/models/result.py

from typing import Any, Dict, List, Optional, Union

import humanfriendly
from pydantic import BaseModel, Field
from pytimeparse.timeparse import timeparse


class LocalFile(BaseModel):
    src: str


class MemoryFile(BaseModel):
    content: str


class PreparedFile(BaseModel):
    file_id: str = Field(..., alias="fileId")


class Collector(BaseModel):
    name: str
    max: int
    pipe: bool = True


class Symlink(BaseModel):
    symlink: str


class StreamIn(BaseModel):
    stream_in: bool = Field(..., alias="streamIn")


class StreamOut(BaseModel):
    stream_out: bool = Field(..., alias="streamOut")


InputFile = Union[LocalFile | MemoryFile | PreparedFile | Symlink]


class Cmd(BaseModel):
    args: Optional[List[str]] = None
    env: List[str] = []
    stdin: Optional[Union[InputFile | StreamIn]] = None
    stdout: Optional[Union[Collector | StreamOut]] = None
    stderr: Optional[Union[Collector | StreamOut]] = None
    cpu_limit: int = Field(0, serialization_alias="cpuLimit")
    real_cpu_limit: int = Field(0, serialization_alias="realCpuLimit")
    clock_limit: int = Field(2 * timeparse("1s"), serialization_alias="clockLimit")
    memory_limit: int = Field(
        humanfriendly.parse_size("128m"), serialization_alias="memoryLimit"
    )
    stack_limit: int = Field(0, serialization_alias="stackLimit")
    proc_limit: int = Field(50, serialization_alias="procLimit")
    cpu_rate_limit: int = Field(0, serialization_alias="cpuRateLimit")
    cpu_set_limit: str = Field("", serialization_alias="cpuSetLimit")
    copy_in: Dict[str, InputFile] = Field({}, serialization_alias="copyIn")
    copy_in_cached: Dict[str, str] = Field({}, serialization_alias="copyInCached")
    copy_in_dir: str = Field(".", serialization_alias="copyInDir")
    # reconsider this default situation
    copy_out: List[str] = Field(["stdout", "stderr"], serialization_alias="copyOut")
    copy_out_cached: List[str] = Field([], serialization_alias="copyOutCached")
    copy_out_max: int = Field(0, serialization_alias="copyOutMax")
    copy_out_dir: str = Field("", serialization_alias="copyOutDir")
    tty: bool = False
    strict_memory_limit: bool = Field(False, serialization_alias="strictMemoryLimit")
    data_segment_limit: bool = Field(False, serialization_alias="dataSegmentLimit")
    address_space_limit: bool = Field(False, serialization_alias="addressSpaceLimit")


class OptionalCmd(BaseModel):
    args: Optional[List[str]] = None
    env: Optional[List[str]] = None
    stdin: Optional[Union[InputFile | StreamIn]] = None
    stdout: Optional[Union[Collector | StreamOut]] = None
    stderr: Optional[Union[Collector | StreamOut]] = None
    cpu_limit: Optional[int] = Field(None, serialization_alias="cpuLimit")
    real_cpu_limit: Optional[int] = Field(None, serialization_alias="realCpuLimit")
    clock_limit: Optional[int] = Field(
        2 * timeparse("1s"), serialization_alias="clockLimit"
    )
    memory_limit: Optional[int] = Field(
        humanfriendly.parse_size("128m"), serialization_alias="memoryLimit"
    )
    stack_limit: Optional[int] = Field(None, serialization_alias="stackLimit")
    proc_limit: Optional[int] = Field(50, serialization_alias="procLimit")
    cpu_rate_limit: Optional[int] = Field(None, serialization_alias="cpuRateLimit")
    cpu_set_limit: Optional[str] = Field(None, serialization_alias="cpuSetLimit")
    copy_in: Optional[Dict[str, InputFile]] = Field(None, serialization_alias="copyIn")
    copy_in_cached: Optional[Dict[str, str]] = Field(
        None, serialization_alias="copyInCached"
    )
    copy_in_dir: Optional[str] = Field(None, serialization_alias="copyInDir")
    copy_out: Optional[List[str]] = Field(
        ["stdout", "stderr"], serialization_alias="copyOut"
    )
    copy_out_cached: Optional[List[str]] = Field(
        None, serialization_alias="copyOutCached"
    )
    copy_out_max: Optional[int] = Field(None, serialization_alias="copyOutMax")
    copy_out_dir: Optional[str] = Field(None, serialization_alias="copyOutDir")
    tty: Optional[bool] = None
    strict_memory_limit: Optional[bool] = Field(
        None, serialization_alias="strictMemoryLimit"
    )
    data_segment_limit: Optional[bool] = Field(
        None, serialization_alias="dataSegmentLimit"
    )
    address_space_limit: Optional[bool] = Field(
        None, serialization_alias="addressSpaceLimit"
    )


class ExecutorWith(BaseModel):
    default: Cmd
    cases: List[OptionalCmd]


class Executor(BaseModel):
    name: str
    with_: ExecutorWith = Field(..., serialization_alias="with")


class ParserConfig(BaseModel):
    name: str
    with_: Dict[str, Any] = Field(..., serialization_alias="with")


class StageDetail(BaseModel):
    name: str
    group: Optional[str] = ""
    executor: Executor
    parsers: List[ParserConfig]


class Stage(BaseModel):
    sandbox_exec_server: str = Field(
        "172.17.0.1:5051", serialization_alias="sandboxExecServer"
    )
    sandbox_token: str = Field("", serialization_alias="sandboxToken")
    output_path: str = Field(
        "/tmp/joj3_result.json", serialization_alias="outputPath"
    )  # nosec: B108
    stages: List[StageDetail]
    pre_stages: List[StageDetail] = Field([], serialization_alias="preStages")
    post_stages: List[StageDetail] = Field([], serialization_alias="postStages")


class Config(BaseModel):
    name: str = ""
    log_path: str = Field("", serialization_alias="logPath")
    expire_unix_timestamp: int = Field(-1, serialization_alias="expireUnixTimestamp")
    effective_unix_timestamp: int = Field(
        -1, serialization_alias="effectiveUnixTimestamp"
    )
    actor_csv_path: str = Field("", serialization_alias="actorCsvPath")
    max_total_score: int = Field(100, serialization_alias="maxTotalScore")
    stage: Stage