drluigi/interface.py at main · xnorkl/drluigi · GitHub

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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
import argparse

import config as conf
config = conf.load()

# TODO move to click instead.


def parser():
    """ Command Line Args """

    # Bottom Menu
    leaf = argparse.ArgumentParser(add_help=False)

    leaf.add_argument(
        '-p', '--payload',
        help='Payload to to be injected into file.'
    )

    leaf.add_argument(
        '--debug',
        action='store_true'
    )

    userprovided = leaf.add_argument_group(title='User provided file options')
    userprovided.add_argument(
        '-f', '--file',
        help='Provide a local file.'
    )

    userprovided.add_argument(
        '-u', '--url',
        help='Provide a url to a file.'
    )

    websearch = leaf.add_argument_group(title='Web search options')
    websearch.add_argument(
        '-q', '--query',
        help='Provide a search query.'
    )

    websearch.add_argument(
        '-s', '--searchmode',
        nargs='?',
        choices=config['searchmode'],
        help='Select a searchmode.'
    )


    # Outputs menu
    outputmenu = argparse.ArgumentParser(add_help=False)
    output = outputmenu.add_argument_group(title='Output options')

    output.add_argument(
        '-o', '--output',
        help='Provide output path.'
    )

    output.add_argument(
        '--json',
        action='store_true',
        help='Pass result as json object.'
    )


    # Submenu
    submenu = argparse.ArgumentParser(add_help=False)

    typemenu = submenu.add_subparsers(
        dest='type'
    )

    image = typemenu.add_parser(
        'image',
        parents=[leaf, outputmenu],
        help='Use image procedures.'
    )

    image.add_argument(
        'image injection',
        const='comment',
        nargs='?',
        choices=config['injection']['image'],
        help='Select an injection method.'
    )

    pdf = typemenu.add_parser(
        'pdf',
        parents=[leaf, outputmenu],
        help='Use pdf injection procedures.'
    )

    rtf = typemenu.add_parser(
        'rtf',
        parents=[leaf, outputmenu],
        help='Use pdf inject procedures.'
    )


    # Main menu
    menu = argparse.ArgumentParser(
        description='Welcome to Dr. Luigi\'s malware clinic.'
    )

    mainmenu = menu.add_subparsers(
        dest='operation'
    )

    mainmenu.add_parser(
        'inject',
        parents=[submenu],
        help='Use injection options.'
    )

    mainmenu.add_parser(
        'analyse',
        parents=[submenu],
        help='Use analysis options.'
    )

    menu.set_defaults(usage=lambda n: menu.print_usage())
    menu.set_defaults(help=lambda n: menu.print_help())

    args = menu.parse_args()

    # Custom Error handling.

    # Currently you cannot nest objects in mutually exclusive groups.
    userprovided_opts = args.url or args.file
    websearch_opts = args.query or args.searchmode


    # There may be a way to remove a positional argument from the help menu.'
    positional_args = args._get_kwargs()[-1][-1]
    if args.operation == 'analyse'and positional_args:
        menu.error('Analyse option is not compatible with injection options.')

    if userprovided_opts and websearch_opts:
        menu.error('User provided options are incompatible with web search options.')
    else:
        return args