-
Notifications
You must be signed in to change notification settings - Fork 7
/
cli.py
88 lines (67 loc) · 2.81 KB
/
cli.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
import click
import re
from loguru import logger
from utils import get_template_from_token, get_tainted_packets, get_tainted_packet, get_clean_source_packets
from stub import stub
from drivers import MockDriver
@click.group()
@click.pass_context
def cli(ctx):
"""
\b
__________ ________ ____ __.
\______ \ _______ _\_____ \ | |/ _|
| _// __ \ \/ // | \| <
| | \ ___/\ // | \ | \
|____|_ /\___ >\_/ \_______ /____|__ \\
\/ \/ \/ \/
"""
pass
@cli.command()
@click.argument('packets_log', type=click.File('r'))
@click.argument('scanner_report', type=click.File('r'))
@click.option('--token-format', required=False)
@click.pass_context
def taint(ctx, packets_log, scanner_report, token_format):
"""
Get tainted paths between log and report.
"""
token_format_regex = re.compile(r'[\w]{8}-[\w]{4}-[\w]{4}-[\w]{4}-[\w]{12}', re.I)
# token_format_regex = re.compile('[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]-{4}[0-9a-f]{4}-[0-9a-f]{12}\Z', re.I)
if token_format:
token_format_regex = re.compile(token_format, re.I)
source_packets = get_clean_source_packets(packets_log)
scanner_report_content = scanner_report.read()
packets_log_tokens = [ token_format_regex.findall(packet) for packet in source_packets]
packets_log_tokens = [token for sublist in packets_log_tokens for token in sublist]
logger.debug(packets_log_tokens)
scanner_report_tokens = token_format_regex.findall(scanner_report_content)
common_tokens = set(packets_log_tokens) and set(scanner_report_tokens)
# logger.info(common_tokens)
logger.success('{} tainted path(s) found.'.format(len(common_tokens)))
for token in common_tokens:
# logger.success(token)
# TODO: the [2:-1] trick removes the b' and ' from the packet. Find a cleaner way
# TODO: find a better way to convert to escaped string
tainted_source_packets = get_tainted_packets(source_packets, token)
logger.debug(tainted_source_packets)
for packet in tainted_source_packets:
logger.success("Tainted source packet for token {}\n{}".format(token, packet))
@cli.command()
@click.argument('packets_log', type=click.File('r'))
@click.argument('token')
@click.option('--placeholder', '-p', default='$a')
@click.pass_context
def template(ctx, packets_log, token, placeholder):
"""
Get a prebuilt template from a tainted token.
The template can be used from the stub component.
WARNING! Use the same placeholder used by the stub.
"""
source_packets = get_clean_source_packets(packets_log)
template = get_template_from_token(source_packets, token, placeholder)
# logger.debug(template)
click.echo(template)
cli.add_command(stub)
if __name__ == '__main__':
cli()