current community

your communities

more stack exchange communities

company blog
Stack Exchange Inbox Reputation and Badges
tour help
Code Review Stack Exchange is a question and answer site for peer programmer code reviews. Join them; it only takes a minute:

Sign up
Here's how it works:
  1. Anybody can ask a question
  2. Anybody can answer
  3. The best answers are voted up and rise to the top
up vote 3 down vote favorite
2

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?

share|improve this question
1  
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. – Mast 9 hours ago
up vote 8 down vote accepted

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)
share|improve this answer
    
IMO ArgParse is really annoying and subpar compared to docopt but maybe I'll add my own answer. – cat 7 hours ago
1  
@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. – Graipher 6 hours ago
    
@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. – Graipher 3 hours ago

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

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