Added regex to make entering IPs directly to work. Also added modifications to make it more in line iwht PEP8

ip-api.py

@@ -1,22 +1,27 @@
 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()
 
-        print('\nCity\State: {}, {}\n Country:   {}\n ISP:       {}\n IP:        {}\n MX:        {}'.format(
+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'],
-            host))
+               x))
+
 
 def findMX(host):
     p = os.popen('host -t MX ' + host)
@@ -36,13 +41,18 @@ def findMX(host):
     # 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)