4
\$\begingroup\$

I've been writing a basic script to enumerate SMTP users (via a user dictionary) on poorly configured SMTP servers. In scripts like this, I usually see arguments handled as follows:

if (len(sys.argv) != 3:
    print "Usage: ..."
    sys.exit(0)

I found the if and print statement approach irritating, but without a concrete reason why in my mind. I decided to use assertions and combine argument handling with general setup. The following is the complete script.

#!/usr/bin/python

import socket
import sys
import os

SMTP_PORT = 25

#
# Helper Functions
#


def create_connection(host, port):
    smtp_server = sys.argv[1]
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect((host, port))

    return(sock)


def smtp_verify(username, smtp_conn):
    smtp_conn.send('VRFY ' + username + '\r\n')
    smtp_reply = smtp_conn.recv(1024)

    return(smtp_reply)

#
# Arg checking, connection setup
#


try:
    assert len(sys.argv) == 3, \
    "Usage: smtp_enum.py <smtp_server_ip> <username_list_file>"

    users_file = sys.argv[2]
    assert os.path.isfile(users_file), "Cannot open file %s" % users_file

    smtp_server = sys.argv[1]
    smtp_conn = create_connection(smtp_server, SMTP_PORT)

    banner = smtp_conn.recv(1024)
    print banner
except Exception as e:
    print "Error: %s" % e
    exit(1)

#
# Perform SMTP enumeration
#


with open(users_file, 'r') as users:
    for user in users:
        print smtp_verify(user.strip(), smtp_conn)

smtp_conn.close()

I'm wondering if I may later regret the assertion approach if the script grows. Is this an abuse of assertions, or somehow problematic?

\$\endgroup\$
  • 1
    \$\begingroup\$ I have rolled back the last edit. Please see What to do when someone answers. Basically, don't edit the code after answers have been given. \$\endgroup\$ – Mast Dec 5 '16 at 11:18
10
\$\begingroup\$

What you really want is a robust argument parsing. This calls for argparse:

import argparse
import os.path

def is_file(f):
    if os.path.isfile(f):
        return f
    raise OSError("{} does not exist".format(f))

parser = argparse.ArgumentParser(description="SMTP user checker")
parser.add_argument('server', help="IP address of server")
parser.add_argument('usernames', type=is_file, help="Path to list of user names")
parser.add_argument('--port', '-p', type=int, default=25, help="SMTP port of server")
args = parser.parse_args()

smtp_conn = create_connection(args.server, args.port)
print smtp_conn.recv(1024)

Note that it is already standard practice of Python to print any error and exit afterwards.

Additionally, you probably want to close your SMTP connection even in the case of an exception. For this, you could use a contextmanager:

from contextlib import contextmanager


@contextmanager
def create_connection(host, port):
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((host, port))
        yield sock
    except socket.error:
        # Handle specific exception here, or not and re-raise
        raise
    finally:
        # Ensure sock is closed under all circumstances
        try:
            sock.close()
        except UnboundLocalError:
            # Did not even succeed in creating the variable, nothing left to do
            pass

And then use it like this:

with create_connection(args.server, args.port) as smtp_conn:
    print smtp_conn.recv(1024)
    with open(args.usernames) as users:
        for user in users:
            print smtp_verify(user.strip(), smtp_conn)

Note that Python uses read mode by default when opening a file, so there is no need to use 'r'.

Final code:

import argparse
import socket
import os.path
from contextlib import contextmanager


@contextmanager
def create_connection(host, port):
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((host, port))
        yield sock
    except socket.error:
        # Handle specific exception here, or not and re-raise
        raise
    finally:
        # Ensure sock is closed under all circumstances
        try:
            sock.close()
        except UnboundLocalError:
            # Did not even succeed in creating the variable, nothing left to do
            pass


def smtp_verify(username, smtp_conn):
    smtp_conn.send('VRFY {}\r\n'.format(username))
    return smtp_conn.recv(1024)


def is_file(f):
    if os.path.isfile(f):
        return f
    raise OSError("{} does not exist".format(f))


def parse_args():
    parser = argparse.ArgumentParser(description="SMTP user checker")
    parser.add_argument('server', help="IP address of server")
    parser.add_argument('usernames', type=is_file, help="Path to list of user names")
    parser.add_argument('--port', '-p', type=int, default=25, help="SMTP port of server")
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    with create_connection(args.server, args.port) as smtp_conn:
        print smtp_conn.recv(1024)
        with open(args.usernames) as users:
            for user in users:
                print smtp_verify(user.strip(), smtp_conn)

Alternatively, you could use the docopt module for this:

"""SMTP user checker.
Usage:
  argparser_comp.py SERVER USERNAMES [(-p <p>| --port=<p>)]
  argparser_comp.py (-h | --help)
Options:
  -h --help          Show this screen.
  -p <p> --port=<p>  SMTP port of server [default: 25].
"""

import socket
import os.path
from contextlib import contextmanager
from docopt import docopt


def is_file(f):
    if os.path.isfile(f):
        return f
    raise OSError("{} does not exist".format(f))


def parse_args():
    args = docopt(__doc__)
    return args["SERVER"], is_file(args["USERNAMES"]), int(args["--port"])


@contextmanager
def create_connection(host, port):
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((host, port))
        yield sock
    except socket.error:
        # Handle specific exception here, or not and re-raise
        raise
    finally:
        # Ensure sock is closed under all circumstances
        try:
            sock.close()
        except UnboundLocalError:
            # Did not even succeed in creating the variable, nothing left to do
            pass


def smtp_verify(username, smtp_conn):
    smtp_conn.send('VRFY {}\r\n'.format(username))
    return smtp_conn.recv(1024)


if __name__ == "__main__":
    server, usernames, port = parse_args()
    with create_connection(server, port) as smtp_conn:
        print smtp_conn.recv(1024)
        with open(usernames) as users:
            for user in users:
                print smtp_verify(user.strip(), smtp_conn)
\$\endgroup\$
  • \$\begingroup\$ IMO ArgParse is really annoying and subpar compared to docopt but maybe I'll add my own answer. \$\endgroup\$ – cat Dec 5 '16 at 13:34
  • 1
    \$\begingroup\$ @cat You are very welcome to write a docopt answer. Would be nice to compare which is easier to write/understand for such a simple case as this. \$\endgroup\$ – Graipher Dec 5 '16 at 14:00
  • 1
    \$\begingroup\$ @cat I tried creating a docopt string for this. In my opinion it is not really clearer, especially because it lacks type handling and automatic callbacks (to check whether the file exists for example). Have a look at the gist I created for this. \$\endgroup\$ – Graipher Dec 5 '16 at 16:54

Your Answer

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