Allow to specify config name on command line.

And thus support multiple configs - to streamline testing, etc. Config
is passed down to iploc.py .

Change-Id: I07bb511fbe80185a422b771f99dd1dc544097b46

1 files changed, 23 insertions, 8 deletions

diff --git a/iploc.py b/iploc.py
index 602979b..b70c203 100644
--- a/iploc.py
+++ b/iploc.py
@@ -8,9 +8,11 @@ import re
 import struct
 import sys
 import os
+import optparse
 
 from bsddb3 import db
 
+
 def load_config(fname):
     d = {}
     f = open(fname)
@@ -34,15 +36,12 @@ def get_path(fname):
     return fname
 
 mypath = os.path.dirname(sys.argv[0])
-config = load_config(mypath + "/config")
-
-# Setup DB
-IP2LocObj = IP2Location.IP2Location()
-IP2LocObj.open(get_path(config["LOC_DB_FILE"]))
+if not mypath:
+    mypath = "."
+config_name = os.path.join(mypath, "config")
+config = None
 
 REVERSE_DNS_DB = db.DB()
-REVERSE_DNS_DB.open(get_path(config["DNSHISTORY_DB"]), dbtype=db.DB_BTREE)
-atexit.register(REVERSE_DNS_DB.close)
 
 temp_ident = ""
 temp_user = ""
@@ -76,6 +75,14 @@ def get_reverse_dns(ip_address):
 
 
 def main(file_names):
+    global config
+    config = load_config(config_name)
+
+    IP2LocObj = IP2Location.IP2Location()
+    IP2LocObj.open(get_path(config["LOC_DB_FILE"]))
+
+    REVERSE_DNS_DB.open(get_path(config["DNSHISTORY_DB"]), dbtype=db.DB_BTREE)
+    atexit.register(REVERSE_DNS_DB.close)
 
     for line in fileinput.input([file_names]):
         locations = []
@@ -119,5 +126,13 @@ def main(file_names):
 
         del locations
 
+
 if __name__ == '__main__':
-    main(sys.argv[1])
+    from optparse import OptionParser
+    optparser = OptionParser()
+    optparser.add_option("--config",
+                      help="Use this config file", metavar="FILE")
+    options, args = optparser.parse_args()
+    if options.config:
+        config_name = options.config
+    main(args[0])