From c1cc65d4aef76518fdcd206b4862a0f1fa40d4a2 Mon Sep 17 00:00:00 2001
From: Michael Heidelberger <mheide@localhost.localdomain>
Date: Thu, 19 Apr 2018 10:04:44 -0500
Subject: [PATCH] Added regex to make entering IPs directly to work. Also added
 modifications to make it more in line iwht PEP8

---
 ip-api.py | 52 +++++++++++++++++++++++++++++++---------------------
 1 file changed, 31 insertions(+), 21 deletions(-)

diff --git a/ip-api.py b/ip-api.py
index 6080f40..452471f 100644
--- a/ip-api.py
+++ b/ip-api.py
@@ -1,48 +1,58 @@
 import argparse
-import json
 import requests
 import os
+import re
 
-def search(host):
-    main_api = 'http://ip-api.com/json/'
 
-    ip = findMX(host)
-    for host in ip:
-        json_data = requests.get(main_api + host).json()
+def query_api(host):
+    main_api = 'http://ip-api.com/json/'
+    # For every host do an API request
+    for x in host:
+        json_data = requests.get(main_api + x).json()
+        # Print out wanted JSON data formatted nicely
+        print('\nCity\State: {}, {}\n'
+              'Country:   {}\n'
+              'ISP:       {}\n'
+              'IP:        {}\n'
+              'MX:        {}'.format(
+               json_data['city'],
+               json_data['regionName'],
+               json_data['country'],
+               json_data['isp'],
+               json_data['query'],
+               x))
 
-        print('\nCity\State: {}, {}\n Country:   {}\n ISP:       {}\n IP:        {}\n MX:        {}'.format(
-            json_data['city'],
-            json_data['regionName'],
-            json_data['country'],
-            json_data['isp'],
-            json_data['query'],
-            host))
 
 def findMX(host):
     p = os.popen('host -t MX ' + host)
 
-    #initialize dicts
+    # initialize dicts
     std_out = []
     split = []
     MXServer = []
 
-    #append terminal output to variable std_out
+    # append terminal output to variable std_out
     for line in p:
         std_out.append(line)
     p.close
 
-    #create iterator
+    # create iterator
     i = 0
-    #split line into dict and return MX servers
+    # split line into dict and return MX servers
     for x in std_out:
         split = std_out[i].split()
-        MXServer.append(split[6])
+        MXServer.append(split[-1])
         i = i + 1
-    return MXServer
+    query_api(MXServer)
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("host", help="hostname to lookip")
     args = parser.parse_args()
-
-    search(args.host)
+    # If the arg is an IP address go straight to API
+    if re.fullmatch(args.host,
+                    '[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}'):
+        query_api(args.host)
+    else:
+        findMX(args.host)