# -*- coding: utf-8 -*-
# Copyright © 2023 Contrast Security, Inc.
# See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
from contrast.agent.policy.registry import register_propagation_nodes
from contrast.agent.policy.utils import CompositeNode


string_propagators = [
    CompositeNode(
        {
            "module": "BUILTIN",
            "class_name": "str",
            "policy_patch": False,
            "instance_method": True,
            "source": "OBJ",  # This is the default, some specific nodes may override it
            "target": "RETURN",
        },
        [
            {
                # There is no kwarg corresponding to ARG_0
                "method_name": "concat",
                "source": "OBJ,ARG_0",
                "action": "APPEND",
            },
            {
                # This is string multiplication
                # TODO: PYT-1345 SPLAT makes this prone to FPs
                "method_name": "repeat",
                "action": "SPLAT",
            },
            {
                "method_name": "translate",
                "action": "SPLAT",
            },
            {
                "method_name": [
                    "capitalize",
                    "casefold",
                    "lower",
                    "swapcase",
                    "title",
                    "upper",
                    "encode",
                    "decode",
                ],
                "action": "KEEP",
            },
            {
                "method_name": "ljust",
                "action": "APPEND",
            },
            {
                "method_name": ["zfill", "rjust"],
                "action": "PREPEND",
            },
            {
                "method_name": "center",
                "action": "CENTER",
            },
            {
                # replace() takes no keyword arguments
                "method_name": "replace",
                "source": "OBJ,ARG_1",
                "action": "REPLACE",
            },
            {
                "method_name": "format",
                "source": "OBJ,ALL_ARGS,ALL_KWARGS",
                "action": "FORMAT",
            },
            {
                # cformat uses the % operator, so there are no kwargs
                "method_name": "cformat",
                "source": "OBJ,ALL_ARGS",
                "action": "FORMAT",
            },
            {
                # format_map() takes no keyword arguments
                "method_name": "formatmap",
                "source": "OBJ,ALL_ARGS",
                "action": "FORMAT",
            },
            {
                # fstring formatting is accomplished with string joins, and doesn't use kwargs
                "method_name": "fstring",
                "source": "OBJ,ARG_0",
                "action": "JOIN",
            },
            {
                "method_name": [
                    "strip",
                    "lstrip",
                    "rstrip",
                    "removeprefix",
                    "removesuffix",
                ],
                "action": "REMOVE",
            },
            {
                "method_name": ["subscript", "__getitem__"],
                "action": "SLICE",
            },
            {
                # join() takes no keyword arguments
                "method_name": "join",
                "source": "OBJ,ARG_0",
                "action": "JOIN",
            },
            {
                "method_name": [
                    "split",
                    "splitlines",
                    "rsplit",
                    "partition",
                    "rpartition",
                ],
                "action": "SPLIT",
            },
            {
                # This is str(), bytes(), unicode(), or bytearray().
                # The reported name is slightly misleading
                "instance_method": False,
                "method_name": "CAST",
                "source": "ARG_0,KWARG:object",
                "action": "SPLAT",
            },
            {
                "method_name": "__repr__",
                "action": "REPR",
            },
        ],
    ),
    CompositeNode(
        {
            "module": "operator",
            "target": "RETURN",
            "policy_patch": False,
        },
        [
            {
                "method_name": ["add", "iadd"],
                "source": "ARG_0,ARG_1",
                "action": "APPEND",
            },
            {
                "method_name": "mod",
                "source": "ARG_0,ARG_1",
                "action": "FORMAT",
            },
        ],
    ),
    {
        "module": "BUILTIN",
        "method_name": ["ascii", "format"],
        # No keyword arguments
        "source": "ARG_0",
        "target": "RETURN",
        "action": "REPR",
    },
]


register_propagation_nodes(string_propagators)
