Verifying user accounts on SMTP servers

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  "

    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?

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 | --port=)]
  argparser_comp.py (-h | --help)
Options:
  -h --help          Show this screen.
  -p  --port=  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)

StackExchange Code Review Q#148984, answer score: 9