some tweaks

1 files changed, 11 insertions, 8 deletions

diff --git a/tor-site-tester.py b/tor-site-tester.py
index ac1d112..912117d 100644
--- a/tor-site-tester.py
+++ b/tor-site-tester.py
@@ -180,22 +180,25 @@ def table(tbl):
         for i, cell in enumerate(row):
             w[i] = max(len(str(cell)), w[i])
 
-    print("+" + "+".join("-" * (x+2) for x in w) + "+")
+    divider = ("+" + "+".join("-" * (x+2) for x in w) + "+\n")
+    tablestr = ""
+    tablestr += divider
     for j, row in enumerate(tbl):
         for i, cell in enumerate(row):
             c = cell
             if c == "200": c = f"\u001b[32m{c}\u001b[0m"
             if c == "403": c = f"\u001b[31m{c}\u001b[0m"
-            print(f"| {c:{w[i]}} ", end="")
-        print("|")
+            tablestr += f"| {c:{w[i]}} "
+        tablestr += "|\n"
         if j == 0:
-            print("+" + "+".join("-" * (x+2) for x in w) + "+")
-    print("+" + "+".join("-" * (x+2) for x in w) + "+")
+            tablestr += divider
+    tablestr += divider
+    Term.log(tablestr)
 
 if __name__ == "__main__":
     parser = ArgumentParser()
     parser.add_argument("--port-start", type=int, default=7000)
-    parser.add_argument("-c", "--countries", type=str, required=True)
+    parser.add_argument("-c", "--countries", type=str, default="")
     parser.add_argument("-s", "--sites-file", type=str, required=True)
     parser.add_argument("-o", "--out-file", type=str, default=None)
     args = parser.parse_args()
@@ -205,11 +208,11 @@ if __name__ == "__main__":
     Term.proc_start("Read sites", spin=True)
     with open(args.sites_file, "r") as f:
         sites = [x.strip() for x in f.readlines() if x.strip() != ""]
-    Term.proc_end()
+    Term.proc_end(msg=f"got {len(sites)} to test")
 
     tor = TORConnections(args.port_start)
     try:
-        countries = args.countries.split(",")
+        countries = [x.strip() for x in args.countries.split(",") if x.strip() != ""]
         results = test_sites(tor, sites, countries)
         pp_results(countries, results, args.out_file)
     except KeyboardInterrupt: