3
\$\begingroup\$

I'm working on the OSCP, and Information Security certificate based around penetration testing. CLI-based tools seem to be the core of most pen-testing assignments, so I'd like to build a tool aiding in command/switch searching.

The goal of the program will to (eventually) load command line arguments by searching in real-time based on arguments, application names, common use-case, and application details. Something like a dynamic Google of the CLI for pen-testing tools.

I'm relatively new to Python. I've never really had anyone look over my code before. I'm hoping that there are not too big of gaps in the logic and syntax. I'm sure there are areas where my code could be simplified. Any suggestions are appreciated.

import threading
import xml.etree.ElementTree as et
import os

# Command to clear the CLI in NT-based systems
cls = lambda: os.system('cls')

# Application dictionary data structure for internal data use
app_dict = {}

# Reference data sets for associative commands and arguments
commands = []
command_key = {}
arguments = []
command_pointers = []
argument_pointers = []


# Function to return query results in a consistent format
def return_command(command_pointer):
    raw_input('''
Application Name: {}
Arguments: {}
    '''.format(command_pointer.name, command_pointer.arguments))


# Base class for a command, identified as the first argument to the CLI
class Command:
    def __init__(self, name):
        self.name = name
        self.arguments = []
        commands.append(name)
        command_key[name] = self
        command_pointers.append(self)

    def set_argument(self, number, optional, prefix):
        self.arguments.append(number)
        self.argument = Argument(self.name, number)
        self.argument.create_argument(optional, prefix)

    def set_argument_details(self, details, common):
        self.argument.set_argument_details(details, common)


# Subsequent CLI arguments following the base command
class Argument:
    def __init__(self, parent_command, number):
        self.parent_command = parent_command
        self.number = number
        arguments.append([parent_command, number])
        argument_pointers.append(self)

    def create_argument(self, optional, prefix):
        self.optional = optional
        self.prefix = prefix

    def set_argument_details(self, details, common):
        self.details = details
        self.common = common


# Starts the XML parsing thread
class ReadThread (threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        xmldoc = et.parse('terms.xml')
        data = xmldoc.getroot()
        for app in data:
            for details in app.iter():
                try:
                    current_command = Command(details.attrib['name'])
                except:
                    current_command.set_argument(details.attrib['number'],
                                                 details.attrib['optional'],
                                                 details.attrib['prefix'])
                    try:
                        current_command.set_argument_details(details.attrib['details'],
                                                             details.attrib['common'])
                    except:
                        pass  # No argument details specified


# Searches XMl loaded commands
# Eventually switching to real-time search ability
class SearchThread (threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)

    def run(self):
        while True:
            cls()
            search_key = raw_input('>>> ')
            if search_key in commands:
                return_command(command_key[search_key])


# Program Start
ReadThread()
query = SearchThread()
query.run()

Here is an example of the base XML that is loaded on program start. I am using this data source until I can implement a functioning DB.

<?xml version="1.0" encoding="UTF-8"?>
<data>
<app name="Test Application 1">
<arg number="1" optional="no" prefix=""></arg>
<arg number="2" optional="yes" prefix="-" details="This is the second argument to the command, and is optional." common="Common use syntax"></arg>  
</app>
<app name="Test Application 2">
<arg number="1" optional="no" prefix=""></arg>
<arg number="2" optional="yes" prefix="-" details="This is the second argument to the command, and is optional." common="Common use syntax"></arg>  
</data>

It's working, so far. Here is an example of output searching the whole string of "Test Application 2". Eventually, a dynamic search would be added that would parse the input string dynamically and return results in real time against more than just the application name.

>>> Test Application 2

Application Name: Test Application 2
Arguments: ['1', '2']
\$\endgroup\$
3
\$\begingroup\$

Coding style

There are many coding style issues here. PEP8 is the official Python coding style guide, give it a good read and start following it. There is a command line tool named pep8 which you can install easily with pip install pep8 and run on your code. Try to fix all reported problems.

Common Python practices.

Classes should have CamelCase names, so Command instead of command and Argument instead of argument.

Classes should extend object, so class Command(object): instead of class command:.

You should define all instance attributes in the constructor. For example this is not so good:

def __init__(self, name):
    self.name = name

def set_argument(self, number):
    self.argument = Argument(self.name, number)

You should initialize self.argument in __init__:

def __init__(self, name):
    self.name = name
    self.argument = None

def set_argument(self, number):
    self.argument = Argument(self.name, number)

Class and method decomposition

The way you divided the problem to classes and methods is very strange, confusing and fragile. Consider how you create a Command class:

  • Create an empty Command
  • Set an argument in Command
    • Creates an Argument
    • Set attributes on Argument
  • Set argument details in Command
    • Delegates to Argument

This is really twisted logic. You are piecing together a Command object like a jigsaw puzzle, sticking part by part on it. Ideally, it's best to give a class everything it needs to function correctly in the constructor. You could refactor your code to behave more like this:

  • Create a collection of Arguments that will be part of a Command
  • Create the Command with the collection of Arguments

Something like this would be better:

commands = {}


def return_command(command):
    raw_input(textwrap.dedent('''
        Application Name: {}
        Arguments: {}
        '''.format(command.name, [arg.number for arg in command.arguments])))


class Command:
    def __init__(self, name, arguments):
        self.name = name
        self.arguments = arguments


class Argument:
    def __init__(self, number, optional, prefix, details, common):
        self.number = number
        self.optional = optional
        self.prefix = prefix
        self.details = details
        self.common = common


def parse_arguments(app):
    for details in app:
        argument = Argument(details.attrib['number'], details.attrib['optional'],
                            details.attrib['prefix'],
                            details.attrib.get('details'), details.attrib.get('common'))
        yield argument


def parse_xml(path):
    xmldoc = et.parse(path)
    data = xmldoc.getroot()
    for app in data:
        arguments = [_ for _ in parse_arguments(app)]
        command = Command(app.attrib['name'], arguments)
        commands[command.name] = command


class Search(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)

    def run(self):
        while True:
            cls()
            search_key = raw_input('>>> ')
            if search_key in commands:
                return_command(commands[search_key])


parse_xml('terms.xml')
query = Search()
query.run()

That is, the classes should know much about the "outside world". Command should not know more than necessary about an Argument and vice versa. Most important of all, a class should not modify global variables like you did in your version. Notice that I deleted most global variables, left only commands, which is populated and used outside of the Commands class. I strongly recommend you this article:

http://en.wikipedia.org/wiki/Single_responsibility_principle

\$\endgroup\$
  • \$\begingroup\$ Thank you for providing me with a detailed review. I was unaware that there was a PEP8 style module. Although I have looked at PEP8 before, I've never really used it as a baseline (that will change). It looks like I have quite a bit to work on. Thanks again. \$\endgroup\$ – vbiqvitovs Aug 21 '14 at 13:44
  • \$\begingroup\$ On another note, my goal is to have a dynamic-data design which allows for addition of switches (args). In thinking about the current design of the command creation -> argument creation as a sub-command class, I think this design would hold given that arguments are ever-changing and being added. It wouldn't make sense to recreate the Command class every time a new argument was presented, no? \$\endgroup\$ – vbiqvitovs Aug 21 '14 at 14:56
  • \$\begingroup\$ Maybe I'm misunderstanding, but in your example, arguments are required in order to instantiate a command, are they not? I was hoping to have arguments exist as separate objects that could be referenced back to the original command. \$\endgroup\$ – vbiqvitovs Aug 21 '14 at 15:23
  • \$\begingroup\$ My focus will be to integrate this into a database for dynamic, real-time referencing. If I were to utilize the construction of the Command class including arguments, every time a new argument was added to a command (via a database), the command would need to be reinstantiated. \$\endgroup\$ – vbiqvitovs Aug 21 '14 at 15:29
  • \$\begingroup\$ I see. This was not visible in your original post. To handle that, you can add a add_argument method to the Command class. The method can take an Argument object and append it to self.arguments. \$\endgroup\$ – Stop ongoing harm to Monica Aug 21 '14 at 16:16

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.