Root cause
AI change
Support HTTP URL as an input (#365)
Loading…
How AI contributed
Flawed AI-written codeRoot cause
Support HTTP URL as an input (#365)
Fix
The minimum fix commit below closes the same vulnerable path.
Mixed fix set: + unmarked commit ·
Code comparison
@@ -18,6 +18,7 @@ Type, Union, )+from urllib.parse import ParseResult from warnings import warn from pydantic import BaseModel, Field, root_validator, validator@@ -26,12 +27,13 @@ InvalidClassNameError, cached_property, load_yaml,+ load_yaml_from_path, snooper_to_methods, ) from datamodel_code_generator.format import PythonVersion from datamodel_code_generator.model import DataModel, DataModelFieldBase from datamodel_code_generator.model.enum import Enum-from datamodel_code_generator.parser import LiteralType+from datamodel_code_generator.parser import DefaultPutDict, LiteralType from ..model import pydantic as pydantic_model from ..parser.base import Parser@@ -225,7 +227,7 @@ def ref_type(self) -> Optional[JSONReference]: class JsonSchemaParser(Parser): def __init__( self,- source: Union[str, Path, List[Path]],+ source: Union[str, Path, List[Path], ParseResult], *, data_model_type: Type[DataModel] = pydantic_model.BaseModel, data_model_root_type: Type[DataModel] = pydantic_model.CustomRootType,@@ -255,6 +257,7 @@ def __init__( strict_nullable: bool = False, use_generic_container_types: bool = False, enable_faux_immutability: bool = False,+ remote_text_cache: Optional[DefaultPutDict[str, str]] = None, ): super().__init__( source=source,@@ -286,9 +289,10 @@ def __init__( strict_nullable=strict_nullable, use_generic_container_types=use_generic_container_types, enable_faux_immutability=enable_faux_immutability,+ remote_text_cache=remote_text_cache, ) - self.remote_object_cache: Dict[str, Dict[str, Any]] = {}+ self.remote_object_cache: DefaultPutDict[str, Dict[str, Any]] = DefaultPutDict() self.raw_obj: Dict[Any, Any] = {} self._root_id: Optional[str] = None self._root_id_base_path: Optional[str] = None@@ -893,44 +897,20 @@ def _get_ref_body(self, resolved_ref: str) -> Dict[Any, Any]: return self._get_ref_body_from_remote(resolved_ref) def _get_ref_body_from_url(self, ref: str) -> Dict[Any, Any]:- if ref[-1] == '#':- ref = ref[:-1]- cached_ref_body: Optional[Dict[str, Any]] = self.remote_object_cache.get(ref)- if cached_ref_body:- return cached_ref_body- # URL Reference – $ref: 'http://path/to/your/resource' Uses the whole document located on the different server.- try:- import httpx- except ImportError: # pragma: no cover- raise Exception(- f'Please run $pip install datamodel-code-generator[http] to resolve URL Reference ref={ref}'- )- raw_body: str = httpx.get(ref).text- # yaml loader can parse json data.- ref_body = load_yaml(raw_body)- self.remote_object_cache[ref] = ref_body- return ref_body+ return self.remote_object_cache.get_or_put(+ ref, default_factory=lambda key: load_yaml(self._get_text_from_url(key))+ ) def _get_ref_body_from_remote(self, resolved_ref: str) -> Dict[Any, Any]: # Remote Reference – $ref: 'document.json' Uses the whole document located on the same server and in # the same location. TODO treat edge case- if resolved_ref[-1] == '#':- resolved_ref = resolved_ref[:-1] full_path = self.base_path / resolved_ref- ref_full_path = str(full_path) - cached_ref_body: Optional[Dict[str, Any]] = self.remote_object_cache.get(- ref_full_path+ return self.remote_object_cache.get_or_put(+ str(full_path),$ref schema / IP CWE-918 SSRFhttp.py get_body() httpx.get(url) schema IP schema $ref 169.254.169.254 localhostsink=httpx.get(url)source=JSON Schema $ref URLguard= IP _validate_url_
@@ -13,7 +13,8 @@ from enum import IntEnum from io import TextIOBase from pathlib import Path-from typing import Any, DefaultDict, Dict, Optional, Sequence, cast+from typing import Any, DefaultDict, Dict, Optional, Sequence, Union, cast+from urllib.parse import ParseResult, urlparse import argcomplete import black@@ -29,6 +30,7 @@ generate, ) from datamodel_code_generator.parser import LiteralType+from datamodel_code_generator.reference import is_url from .format import PythonVersion, is_supported_in_black @@ -53,12 +55,16 @@ def sig_int_handler(_: int, __: Any) -> None: # pragma: no cover arg_parser.add_argument( '--input', help='Input file/directory (default: stdin)', )+arg_parser.add_argument(+ '--url', help='Input file URL. `--input` is ignore when `--url` is used',+) arg_parser.add_argument( '--input-file-type', help='Input file type (default: auto)', choices=[i.value for i in InputFileType], ) arg_parser.add_argument('--output', help='Output file (default: stdout)')+ arg_parser.add_argument( '--base-class', help='Base Class (default: pydantic.BaseModel)', type=str, )@@ -220,6 +226,16 @@ def validate_path(cls, value: Any) -> Optional[Path]: return value # pragma: no cover return Path(value).expanduser().resolve() + @validator('url', pre=True)+ def validate_url(cls, value: Any) -> Optional[ParseResult]:+ if isinstance(value, str) and is_url(value): # pragma: no cover+ return urlparse(value)+ elif value is None: # pragma: no cover+ return None+ raise Error(+ f'This protocol doesn\'t support only http/https. --input={value}'+ ) # pragma: no cover+ @root_validator def validate_literal_option(cls, values: Dict[str, Any]) -> Dict[str, Any]: if values.get('enum_field_as_literal'):@@ -244,7 +260,7 @@ def validate_use_generic_container_types( ) return values - input: Optional[Path]+ input: Optional[Union[Path, str]] input_file_type: InputFileType = InputFileType.Auto output: Optional[Path] debug: bool = False@@ -271,6 +287,7 @@ def validate_use_generic_container_types( strict_nullable: bool = False use_generic_container_types: bool = False enable_faux_immutability: bool = False+ url: Optional[ParseResult] = None def merge_args(self, args: Namespace) -> None: for field_name in self.__fields__:@@ -362,7 +379,7 @@ def main(args: Optional[Sequence[str]] = None) -> Exit: try: generate(- input_=config.input or sys.stdin.read(),+ input_=config.url or config.input or sys.stdin.read(), input_file_type=config.input_file_type, output=config.output, target_python_version=config.target_python_version,AI introduced this behavior: `@validator('url', pre=True)`
@@ -7,6 +7,7 @@ from typing import ( TYPE_CHECKING, Any,+ Callable, ClassVar, DefaultDict, Dict,@@ -18,6 +19,7 @@ Sequence, Set, Tuple,+ TypeVar, Union, ) @@ -106,13 +108,28 @@ def short_name(self) -> str: ID_PATTERN: Pattern[str] = re.compile(r'^#[^/].*') +T = TypeVar('T')+++@contextmanager+def context_variable(+ setter: Callable[[T], None], current_value: T, new_value: T+) -> Generator[None, None, None]:+ previous_value: T = current_value+ setter(new_value)+ try:+ yield+ finally:+ setter(previous_value)+ class ModelResolver: def __init__( self, aliases: Optional[Mapping[str, str]] = None, exclude_names: Set[str] = None, duplicate_name_suffix: Optional[str] = None,+ base_url: Optional[str] = None, ) -> None: self.references: Dict[str, Reference] = {} self.aliases: Mapping[str, str] = {} if aliases is None else {**aliases}@@ -122,6 +139,23 @@ def __init__( self.after_load_files: Set[str] = set() self.exclude_names: Set[str] = exclude_names or set() self.duplicate_name_suffix: Optional[str] = duplicate_name_suffix+ self._base_url: Optional[str] = base_url++ @property+ def base_url(self) -> Optional[str]:+ return self._base_url++ def set_base_url(self, base_url: Optional[str]) -> None:+ self._base_url = base_url++ @contextmanager+ def base_url_context(self, base_url: str) -> Generator[None, None, None]:++ if self._base_url:+ with context_variable(self.set_base_url, self.base_url, base_url):+ yield+ else:+ yield @property def current_root(self) -> Sequence[str]:@@ -136,10 +170,8 @@ def set_current_root(self, current_root: Sequence[str]) -> None: def current_root_context( self, current_root: Sequence[str] ) -> Generator[None, None, None]:- previous_root_path: Sequence[str] = self.current_root- self.set_current_root(current_root)- yield- self.set_current_root(previous_root_path)+ with context_variable(self.set_current_root, self.current_root, current_root):+ yield @property def root_id_base_path(self) -> Optional[str]:@@ -157,20 +189,28 @@ def resolve_ref(self, path: Union[Sequence[str], str]) -> str: else: joined_path = self.join_path(path) if ID_PATTERN.match(joined_path):- return self.ids['/'.join(self.current_root)][joined_path]+ ref: str = self.ids['/'.join(self.current_root)][joined_path] elif '#' in joined_path: if joined_path[0] == '#':$ref schema / IP CWE-918 SSRFhttp.py get_body() httpx.get(url) schema IP schema $ref 169.254.169.254 localhostsink=httpx.get(url)source=JSON Schema $ref URLguard= IP _validate_url_
Releases