ronantakizawa/codereview-bench huggingface dataset:
Repository: mandiant/capa
File: capa/render/result_document.py (Python)

Code before:
# Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
# You may obtain a copy of the License at: [package root]/LICENSE.txt
# Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.
import datetime
import collections
from typing import Any, Dict, List, Tuple, Union, Optional

from pydantic import Field, BaseModel

import capa.rules
import capa.engine
import capa.features.common
import capa.features.freeze as frz
import capa.features.address
import capa.features.freeze.features as frzf
from capa.rules import RuleSet
from capa.engine import MatchResults
from capa.helpers import assert_never


class FrozenModel(BaseModel):
    class Config:
        extra = "forbid"


class Sample(FrozenModel):
    md5: str
    sha1: str
    sha256: str
    path: str


class BasicBlockLayout(FrozenModel):
    address: frz.Address


class FunctionLayout(FrozenModel):
    address: frz.Address
    matched_basic_blocks: Tuple[BasicBlockLayout, ...]


class Layout(FrozenModel):
    functions: Tuple[FunctionLayout, ...]


class LibraryFunction(FrozenModel):
    address: frz.Address
    name: str

Review comment (suggestion):
lets not call this `Frozen` if its not frozen.

for things that can still stay frozen, lets use the existing `FrozenModel`. but lets also introduce a `Model` class for things that aren't frozen.

Code after:
# Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
# You may obtain a copy of the License at: [package root]/LICENSE.txt
# Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.
import datetime
import collections
from typing import Any, Dict, List, Tuple, Union, Optional

from pydantic import Field, BaseModel

import capa.rules
import capa.engine
import capa.features.common
import capa.features.freeze as frz
import capa.features.address
import capa.features.freeze.features as frzf
from capa.rules import RuleSet
from capa.engine import MatchResults
from capa.helpers import assert_never


class FrozenModel(BaseModel):
    class Config:
        frozen = True
        extra = "forbid"


class Model(BaseModel):
    class Config:
        extra = "forbid"


class Sample(Model):
    md5: str
    sha1: str
    sha256: str
    path: str


class BasicBlockLayout(Model):
    address: frz.Address


class FunctionLayout(Model):
    address: frz.Address
    matched_basic_blocks: Tuple[BasicBlockLayout, ...]


class Layout(Model):
Repository: mandiant/capa
File: capa/features/extractors/base_extractor.py (Python)

Code before:

    def __init__(self):
        super().__init__()

    @abc.abstractmethod
    def get_processes(self) -> Iterator[ProcessHandle]:
        """
        Yields all the child-processes of a parent one.

        Attributes:
            ph: parent process
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def extract_process_features(self, ph: ProcessHandle) -> Iterator[Tuple[Feature, Address]]:
        """
        Yields all the features of a process. These include:
        - file features of the process' image
        - inter-process injection
        - detected dynamic DLL loading
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def get_threads(self, ph: ProcessHandle) -> Iterator[ProcessHandle]:
        """
        Yields all the threads that a process created.

        Attributes:
            ph: parent process
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def extract_thread_features(self, ph: ProcessHandle, th: ThreadHandle) -> Iterator[Tuple[Feature, Address]]:
        """
        Yields all the features of a thread. These include:
        - sequenced api traces
        - files/registris interacted with
        - network activity
        """
        raise NotImplementedError()

    @abc.abstractclassmethod
    def from_trace(cls, trace: TextIO) -> "DynamicExtractor":
        """
        Most sandboxes provide reports in a serialized text format (i.e. JSON for Cuckoo and CAPE).
        This routine takes a file descriptor of such report (analysis trace) and returns a corresponding DynamicExtractor object.
        """
        raise NotImplementedError()

Review comment (suggestion):
```suggestion
    def get_threads(self, ph: ProcessHandle) -> Iterator[ThreadHandle]:
```

Code after:
    @abc.abstractmethod
    def get_processes(self) -> Iterator[ProcessHandle]:
        """
        Enumerate processes in the trace.
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def extract_process_features(self, ph: ProcessHandle) -> Iterator[Tuple[Feature, Address]]:
        """
        Yields all the features of a process. These include:
        - file features of the process' image
        - inter-process injection
        - detected dynamic DLL loading
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def get_threads(self, ph: ProcessHandle) -> Iterator[ThreadHandle]:
        """
        Enumerate threads in the given process.
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def extract_thread_features(self, ph: ProcessHandle, th: ThreadHandle) -> Iterator[Tuple[Feature, Address]]:
        """
        Yields all the features of a thread. These include:
        - sequenced api traces
        - file/registry interactions
        - network activity
        """
        raise NotImplementedError()
Repository: mandiant/capa
File: capa/features/extractors/base_extractor.py (Python)

Code before:

@dataclass
class ProcessHandle:
    """
    reference to a process extracted by the sandbox.

    Attributes:
        pid: process id
        inner: sandbox-specific data
    """

    pid: int
    inner: Any


@dataclass
class ThreadHandle:
    """
    reference to a thread extracted by the sandbox.

    Attributes:
        tid: thread id
        inner: sandbox-specific data
    """

    tid: int
    inner: Any


class DynamicExtractor(FeatureExtractor):
    """
    DynamicExtractor defines the interface for fetching features from a sandbox' analysis of a sample.

    Features are grouped mainly into threads that alongside their meta-features are also grouped into 
    processes (that also have their own features). Other scopes (such as function and file) may also apply 
    for a specific sandbox.

    This class is not instantiated directly; it is the base class for other implementations.
    """

    def __init__(self):
        super().__init__()

    @abc.abstractmethod
    def get_processes(self) -> Iterator[ProcessHandle]:
        """
        Yields all the child-processes of a parent one.

        Attributes:
            ph: parent process
        """

Review comment (suggestion):
should the thread have a reference to its parent?

i guess the thread handle is always provided alongside a process handle, so this isn't necessary today.

Code after:

@dataclass
class ProcessHandle:
    """
    reference to a process extracted by the sandbox.

    Attributes:
        pid: process id
        inner: sandbox-specific data
    """

    pid: int
    inner: Any


@dataclass
class ThreadHandle:
    """
    reference to a thread extracted by the sandbox.

    Attributes:
        tid: thread id
        inner: sandbox-specific data
    """

    tid: int
    inner: Any


class DynamicExtractor(FeatureExtractor):
    """
    DynamicExtractor defines the interface for fetching features from a sandbox' analysis of a sample.

    Features are grouped mainly into threads that alongside their meta-features are also grouped into 
    processes (that also have their own features). Other scopes (such as function and file) may also apply 
    for a specific sandbox.

    This class is not instantiated directly; it is the base class for other implementations.
    """
    @abc.abstractmethod
    def get_processes(self) -> Iterator[ProcessHandle]:
        """
        Enumerate processes in the trace.
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def extract_process_features(self, ph: ProcessHandle) -> Iterator[Tuple[Feature, Address]]:
        """
        Yields all the features of a process. These include:
        - file features of the process' image
Repository: mandiant/capa
File: capa/features/extractors/base_extractor.py (Python)

Code before:
# Copyright (C) 2020 Mandiant, Inc. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
# You may obtain a copy of the License at: [package root]/LICENSE.txt
# Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.

import abc
import dataclasses
from typing import Any, Dict, Tuple, Union, Iterator, TextIO, BinaryIO
from dataclasses import dataclass

import capa.features.address
from capa.features.common import Feature
from capa.features.address import Address, AbsoluteVirtualAddress

# feature extractors may reference functions, BBs, insns by opaque handle values.
# you can use the `.address` property to get and render the address of the feature.
#
# these handles are only consumed by routines on
# the feature extractor from which they were created.


@dataclass
class FunctionHandle:
    """reference to a function recognized by a feature extractor.

    Attributes:
        address: the address of the function.
        inner: extractor-specific data.
        ctx: a context object for the extractor.
    """

    address: Address
    inner: Any

Review comment (suggestion):
```suggestion
from typing import Any, Dict, Tuple, Union, Iterator
```

Code after:
# Copyright (C) 2020 Mandiant, Inc. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
# You may obtain a copy of the License at: [package root]/LICENSE.txt
# Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.

import abc
import dataclasses
from typing import Any, Dict, Tuple, Union, Iterator
from dataclasses import dataclass

import capa.features.address
from capa.features.common import Feature
from capa.features.address import Address, AbsoluteVirtualAddress

# feature extractors may reference functions, BBs, insns by opaque handle values.
# you can use the `.address` property to get and render the address of the feature.
#
# these handles are only consumed by routines on
# the feature extractor from which they were created.


@dataclass
class FunctionHandle:
    """reference to a function recognized by a feature extractor.

    Attributes:
        address: the address of the function.
        inner: extractor-specific data.
        ctx: a context object for the extractor.
    """

    address: Address
    inner: Any
Repository: mandiant/capa
File: capa/features/extractors/base_extractor.py (Python)

Code before:
    def get_processes(self) -> Iterator[ProcessHandle]:
        """
        Yields all the child-processes of a parent one.

        Attributes:
            ph: parent process
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def extract_process_features(self, ph: ProcessHandle) -> Iterator[Tuple[Feature, Address]]:
        """
        Yields all the features of a process. These include:
        - file features of the process' image
        - inter-process injection
        - detected dynamic DLL loading
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def get_threads(self, ph: ProcessHandle) -> Iterator[ThreadHandle]:
        """
        Yields all the threads that a process created.

        Attributes:
            ph: parent process
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def extract_thread_features(self, ph: ProcessHandle, th: ThreadHandle) -> Iterator[Tuple[Feature, Address]]:
        """
        Yields all the features of a thread. These include:
        - sequenced api traces
        - file/registry interactions
        - network activity
        """
        raise NotImplementedError()


Review comment (suggestion):
```suggestion
        Enumerate threads in the given process.
```

Code after:
        """
        Enumerate processes in the trace.
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def extract_process_features(self, ph: ProcessHandle) -> Iterator[Tuple[Feature, Address]]:
        """
        Yields all the features of a process. These include:
        - file features of the process' image
        - inter-process injection
        - detected dynamic DLL loading
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def get_threads(self, ph: ProcessHandle) -> Iterator[ThreadHandle]:
        """
        Enumerate threads in the given process.
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def extract_thread_features(self, ph: ProcessHandle, th: ThreadHandle) -> Iterator[Tuple[Feature, Address]]:
        """
        Yields all the features of a thread. These include:
        - sequenced api traces
        - file/registry interactions
        - network activity
        """
        raise NotImplementedError()
Repository: mandiant/capa
File: capa/helpers.py (Python)

Code before:

def get_file_taste(sample_path: str) -> bytes:
    if not os.path.exists(sample_path):
        raise IOError(f"sample path {sample_path} does not exist or cannot be accessed")
    with open(sample_path, "rb") as f:
        taste = f.read(8)
    return taste


def is_runtime_ida():
    try:
        import idc
    except ImportError:
        return False
    else:
        return True


def assert_never(value) -> NoReturn:
    assert False, f"Unhandled value: {value} ({type(value).__name__})"


def get_format_from_report(sample: str) -> str:
    with open(sample, "rb") as f:
        report = json.load(f)
    if FORMAT_CAPE.upper() in report.keys():
        return FORMAT_CAPE
    return FORMAT_UNKNOWN


def get_format_from_extension(sample: str) -> str:
    format_ = FORMAT_UNKNOWN
    if sample.endswith(EXTENSIONS_SHELLCODE_32):
        format_ = FORMAT_SC32
    elif sample.endswith(EXTENSIONS_SHELLCODE_64):
        format_ = FORMAT_SC64
    elif sample.endswith(EXTENSIONS_DYNAMIC):
        format_ = get_format_from_report(sample)
    return format_


def get_auto_format(path: str) -> str:
    format_ = get_format(path)
    if format_ == FORMAT_UNKNOWN:
        format_ = get_format_from_extension(path)
    if format_ == FORMAT_UNKNOWN:
        raise UnsupportedFormatError()
    return format_


def get_format(sample: str) -> str:

Review comment (suggestion):
```suggestion
    if "CAPE" in report.keys():
```

Code after:

def get_file_taste(sample_path: str) -> bytes:
    if not os.path.exists(sample_path):
        raise IOError(f"sample path {sample_path} does not exist or cannot be accessed")
    with open(sample_path, "rb") as f:
        taste = f.read(8)
    return taste


def is_runtime_ida():
    try:
        import idc
    except ImportError:
        return False
    else:
        return True


def assert_never(value) -> NoReturn:
    assert False, f"Unhandled value: {value} ({type(value).__name__})"


def get_format_from_report(sample: str) -> str:
    with open(sample, "rb") as f:
        report = json.load(f)
    if "CAPE" in report.keys():
        return FORMAT_CAPE
    return FORMAT_UNKNOWN


def get_format_from_extension(sample: str) -> str:
    format_ = FORMAT_UNKNOWN
    if sample.endswith(EXTENSIONS_SHELLCODE_32):
        format_ = FORMAT_SC32
    elif sample.endswith(EXTENSIONS_SHELLCODE_64):
        format_ = FORMAT_SC64
    elif sample.endswith(EXTENSIONS_DYNAMIC):
        format_ = get_format_from_report(sample)
    return format_


def get_auto_format(path: str) -> str:
    format_ = get_format(path)
    if format_ == FORMAT_UNKNOWN:
        format_ = get_format_from_extension(path)
    if format_ == FORMAT_UNKNOWN:
        raise UnsupportedFormatError()
    return format_


def get_format(sample: str) -> str:
Repository: mandiant/capa
File: capa/rules/__init__.py (Python)

Code before:
    "authors",
    "description",
    "lib",
    "scope",
    "att&ck",
    "mbc",
    "references",
    "examples",
)
# these are meta fields that are internal to capa,
# and added during rule reading/construction.
# they may help use manipulate or index rules,
# but should not be exposed to clients.
HIDDEN_META_KEYS = ("capa/nursery", "capa/path")


class Scope(str, Enum):
    FILE = "file"
    PROCESS = "process"
    FUNCTION = "function"
    BASIC_BLOCK = "basic block"
    INSTRUCTION = "instruction"


FILE_SCOPE = Scope.FILE.value
PROCESS_SCOPE = Scope.PROCESS
FUNCTION_SCOPE = Scope.FUNCTION.value
BASIC_BLOCK_SCOPE = Scope.BASIC_BLOCK.value
INSTRUCTION_SCOPE = Scope.INSTRUCTION.value
# used only to specify supported features per scope.
# not used to validate rules.
GLOBAL_SCOPE = "global"


SUPPORTED_FEATURES: Dict[str, Set] = {
    GLOBAL_SCOPE: {
        # these will be added to other scopes, see below.
        capa.features.common.OS,
        capa.features.common.Arch,
        capa.features.common.Format,
    },
    FILE_SCOPE: {
        capa.features.common.MatchedRule,
        capa.features.file.Export,
        capa.features.file.Import,
        capa.features.file.Section,
        capa.features.file.FunctionName,
        capa.features.common.Characteristic("embedded pe"),
        capa.features.common.String,
        capa.features.common.Class,
        capa.features.common.Namespace,

Review comment (suggestion):
```suggestion
PROCESS_SCOPE = Scope.PROCESS.value
```

Code after:
    "authors",
    "description",
    "lib",
    "scope",
    "att&ck",
    "mbc",
    "references",
    "examples",
)
# these are meta fields that are internal to capa,
# and added during rule reading/construction.
# they may help use manipulate or index rules,
# but should not be exposed to clients.
HIDDEN_META_KEYS = ("capa/nursery", "capa/path")


class Scope(str, Enum):
    FILE = "file"
    PROCESS = "process"
    THREAD = "thread"
    FUNCTION = "function"
    BASIC_BLOCK = "basic block"
    INSTRUCTION = "instruction"


FILE_SCOPE = Scope.FILE.value
PROCESS_SCOPE = Scope.PROCESS.value
THREAD_SCOPE = Scope.THREAD.value
FUNCTION_SCOPE = Scope.FUNCTION.value
BASIC_BLOCK_SCOPE = Scope.BASIC_BLOCK.value
INSTRUCTION_SCOPE = Scope.INSTRUCTION.value
# used only to specify supported features per scope.
# not used to validate rules.
GLOBAL_SCOPE = "global"


SUPPORTED_FEATURES: Dict[str, Set] = {
    GLOBAL_SCOPE: {
        # these will be added to other scopes, see below.
        capa.features.common.OS,
        capa.features.common.Arch,
        capa.features.common.Format,
    },
    FILE_SCOPE: {
        capa.features.common.MatchedRule,
        capa.features.file.Export,
        capa.features.file.Import,
        capa.features.file.Section,
        capa.features.file.FunctionName,
        capa.features.common.Characteristic("embedded pe"),
        capa.features.common.String,
Repository: mandiant/capa
File: capa/features/extractors/ghidra/helpers.py (Python)

Code before:
# Copyright (C) 2020 Mandiant, Inc. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
# You may obtain a copy of the License at: [package root]/LICENSE.txt
# Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.
from typing import Any, Dict, Tuple, Iterator, Optional

import ghidra
from ghidra.program.model.symbol import SymbolType

currentProgram: ghidra.program.database.ProgramDB


def fix_bytes(b: int):
    """Transform signed ints from Java into bytes for Python

    args:
        b: signed int returned from Java processing
    """
    return (b & 0xFF).to_bytes(1, "little")


def find_byte_sequence(seq: bytes) -> Iterator[int]:
    """yield all ea of a given byte sequence

    args:
        seq: bytes to search e.g. b"\x01\x03"
    """
    seqstr = "".join([f"\\x{b:02x}" for b in seq])
    # .add(1) to avoid false positives on regular PE files
    ea = findBytes(currentProgram.getMinAddress().add(1), seqstr, 1, 1)  # type: ignore [name-defined]
    for e in ea:
        yield e


def get_bytes(addr: ghidra.program.model.address.Address, length: int) -> bytes:
    """yield length bytes at addr

    args:

Review comment (suggestion):
```suggestion
def fix_byte(b: int) -> bytes:
```

Code after:
# Copyright (C) 2020 Mandiant, Inc. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
# You may obtain a copy of the License at: [package root]/LICENSE.txt
# Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.
from typing import Any, Dict, Tuple, Iterator, Optional

import ghidra
from ghidra.program.model.symbol import SymbolType

currentProgram: ghidra.program.database.ProgramDB


def fix_byte(b: int) -> bytes:
    """Transform signed ints from Java into bytes for Python

    args:
        b: signed int returned from Java processing
    """
    return (b & 0xFF).to_bytes(1, "little")


def find_byte_sequence(seq: bytes) -> Iterator[int]:
    """yield all ea of a given byte sequence

    args:
        seq: bytes to search e.g. b"\x01\x03"
    """
    seqstr = "".join([f"\\x{b:02x}" for b in seq])
    # .add(1) to avoid false positives on regular PE files
    ea = findBytes(currentProgram.getMinAddress().add(1), seqstr, 1, 1)  # type: ignore [name-defined]
    for e in ea:
        yield e


def get_bytes(addr: ghidra.program.model.address.Address, length: int) -> bytes:
    """yield length bytes at addr

    args: