6
\$\begingroup\$

I wrote a simple Python app that is given an IP address and port and the app then finds the PID listening on that IP:PORT.

I would like to hear constructive criticism regarding the bellow code.

import argparse
import sys
import socket
import subprocess
import os


def is_legal_ip(address):
    ip = address.split(":")
    pieces = ip[0].split('.')
    if len(pieces) != 4: 
        return False
    try: 
        return all(0<=int(p)<256 for p in pieces)
    except ValueError: 
        return False

def parse_arguments():
    if len(sys.argv) > 2 or len(sys.argv) == 1:
        print("Please provide a single argument.")
        exit(1)

    parser = argparse.ArgumentParser("Returns a PID listenning on IP:PORT")
    parser.add_argument('ip_port', metavar='IP:PORT', type=str, nargs=1, help='local address')
    args = parser.parse_args()

    if (is_legal_ip(args.ip_port[0]) == False):
        print("Invalid IP address.")
        exit(1)

    return(args.ip_port[0])

def address_to_pid(ip_port):
    devnull = open(os.devnull, 'w')
    output = subprocess.check_output("netstat -nptl".split(), stderr = devnull).splitlines()
    devnull.close()

    for line in output[2::]:
        raw_line = line.split()
        if raw_line[3] == ip_port:
            if raw_line[6] == "-": 
                return "No PID listening on address"
            else:
                return raw_line[6].split(r"/")[0]


def main():
    ip_port = parse_arguments()

    print address_to_pid(ip_port)



if __name__ == "__main__": main()
\$\endgroup\$
8
\$\begingroup\$

First off, a couple of quick tips:

  • Follow the PEP8 style guide
  • Instead of x == False, use not x, for example in if (is_legal_ip(args.ip_port[0]) == False)
  • Omit unnecessary imports (import socket in this example)
  • Use spaces around operators, for example all(0 <= int(p) < 256 for p in pieces)
  • Break lines after colons, for example in if __name__ == "__main__": main()

Leave the job of argument parsing to argparse

No need to work with sys.args yourself. For example you can simply remove these lines:

if len(sys.argv) > 2 or len(sys.argv) == 1:
    print("Please provide a single argument.")
    exit(1)

The way you specified the ip-port arguments with nargs=1, argparse will take care of validating the number of command line arguments.

Simplify the argument parsing

You can omit nargs=1. But in that case, the value of args.ip_port will not be a list, but a simple string.

You can also omit the type=str, as that's the default.

Your usage would become simpler:

parser.add_argument('ip_port', metavar='IP:PORT', help='local address')
args = parser.parse_args()

if not is_legal_ip(args.ip_port):
    parser.error("Invalid IP address.")

return args.ip_port

Use two arguments instead of one

It's awkward that use ip:port, the splitting at multiple places is just noise. It will be simpler and cleaner to use two arguments instead of one:

parser.add_argument('ip', help='local address')
parser.add_argument('port', type=int, help='local port')
args = parser.parse_args()

if not is_legal_ip(args.ip):
    parser.error("Invalid IP address.")

return args.ip, args.port

Notice that this way argparse can take care of validating the port number.

Then in your main function, you can get the ip address and port pair like this:

ip, port = parse_arguments()

Benefit from the error reporting features of argparse

Instead of this:

if not is_legal_ip(args.ip_port):
    print("Invalid IP address.")
    exit(1)

I suggest to use this way:

if not is_legal_ip(args.ip_port):
    parser.error("Invalid IP address.")

argparse will print a usage help along with the error message, and also exit.

A subtlety here is that argparse will exit with code 2 in this case, which is common in case of invalid parameters of UNIX commands.

Working with file handles

When working with files, use with:

with open(os.devnull, 'w') as devnull:
    output = subprocess.check_output("netstat -nl".split(), stderr=devnull).splitlines()

This makes devnull.close() unnecessary, Python will call it when exiting the block.

\$\endgroup\$
2
\$\begingroup\$

Style

You have some spacing issues, for example, this line:

return all(0<=int(p)<256 for p in pieces)

Should be expanded to this:

return all(0 <= int(p) < 256 for p in pieces)

You also have some unnecessary parentheses. This line:

return(args.ip_port[0])

Should be changed to:

return args.ip_port[0]

This line:

if (is_legal_ip(args.ip_port[0]) == False):

Should also be changed to this:

if is_legal_ip(args.ip_port[0]) == False:

You should also also have two blank lines between functions, not one.

You should also read PEP8, Python's official style guide.


Issues

You should also be using a context manager to open and close files, not file.close. For example, this block of code:

devnull = open(os.devnull, 'w')
output = subprocess.check_output("netstat -nptl".split(), stderr = devnull).splitlines()
devnull.close()

Should be changed to this:

with open(os.devnull, 'w') as devnull:
    output = subprocess.check_output("netstat -nl".split(), stderr=devnull).splitlines()

Miscellaneous nitpicks

This block of code:

for line in output[2::]:
    raw_line = line.split()
    if raw_line[3] == ip_port:
        if raw_line[6] == "-": 
            return "No PID listening on address"
        else:
            return raw_line[6].split(r"/")[0]

Can be shortened to this:

for line in output[2::]:
    raw_line = line.split()
    if raw_line[3] == ip_port and if raw_line[6] == "-": 
        return "No PID listening on address"
    return raw_line[6].split(r"/")[0]

This condition:

if len(sys.argv) > 2 or len(sys.argv) == 1:

Can also be shortened to this:

if len(sys.argv) >= 1:

In addition

\$\endgroup\$
1
\$\begingroup\$

You can try something like this that is pep8 like suggested above and leverages some built-ins from Python.

note: This assumes you want to look for a specific IPv4 NetAddr, and not just stuff listening. For example, the code below will not work with IPv6 addresses (without mods) and will not catch stuff listening on * (such as SSH). This is also python3 friendly (although you may have to adjust some things for psutil to make it work with python3, but the print statements are python3 ready. )

'''
my program does cool stuff
'''
import argparse
import psutil
from ipaddress import IPv4Network, AddressValueError



def is_legal_ip(address):
    '''
    determine if IP is valid
    '''
    address = address[0] 
    try:
        if IPv4Network(unicode(address)):
            pass
    except AddressValueError as valueerr:
        print('invalid IP provided: {} : {}'.format(address, valueerr))
        return False
    return True

def parse_arguments():
    '''
    parse the command line arguments
    '''
    parser = argparse.ArgumentParser('Returns a PID listenning on IP:PORT')
    parser.add_argument('ip_address',
                        type=str,
                        nargs=1,
                        help='local address')
    parser.add_argument('ip_port',
                        type=int,
                        nargs=1,
                        help='local port')
    args = parser.parse_args()
    if not is_legal_ip(args.ip_address):
        parser.error('Invalid IP address : {}'.format(args.ip_address))

    return(args.ip_address, args.ip_port)

def address_to_pid(ip_address='', ip_port=0):
    '''
    determine pid based on IP and port combo
    leverage psutil library and net_connections method
    return the pid of the listening binary
    '''
    possible_pids = set()
    [x.pid for x in psutil.net_connections() if x.laddr == (
        str(ip_address), ip_port) and x.pid not in possible_pids\
                and possible_pids.add(x.pid)]
    if len(possible_pids) < 1:
        return 'Nothing'
    return possible_pids.pop()


def main():
    '''
    a small program to get the PID listening on an IP/port combo
    '''
    ip_address, ip_port = parse_arguments()
    actual_pid = address_to_pid(ip_address[0], ip_port[0])
    print('{} is listening on {} port {}'.format(actual_pid,
                                                 ip_address,
                                                 ip_port))


if __name__ == '__main__':
    main()
\$\endgroup\$

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.