From c52646aebbee3edfc3f23039d676cc596cf06e27 Mon Sep 17 00:00:00 2001
From: Karim Rahal <karim@karimrahal.com>
Date: Tue, 14 Dec 2021 21:08:33 -0800
Subject: [PATCH] Add stdin support

---
 waybacked/__main__.py | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/waybacked/__main__.py b/waybacked/__main__.py
index 2e4879e..4f79e5c 100644
--- a/waybacked/__main__.py
+++ b/waybacked/__main__.py
@@ -1,4 +1,5 @@
 import argparse
+import sys
 import time
 from functools import wraps
 from random import randint
@@ -59,20 +60,29 @@ def run(self):
                 yield data
 
 
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("search", metavar="search", nargs=None, help="Search query")
-    args = parser.parse_args()
-    search = WaybackSearch(args.search)
+def process_search(search):
+    search = WaybackSearch(search)
     try:
         for page in search.run():
             for url in page:
                 print(url, flush=True)
-    except KeyboardInterrupt:
-        pass
     except WaybackSearchError as e:
         print(e)
 
 
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "search", metavar="search", nargs="?", help="Search query", default=None
+    )
+    args = parser.parse_args()
+    search_strs = (args.search,) if args.search else (l.rstrip("\n") for l in sys.stdin)
+    try:
+        for search in search_strs:
+            process_search(search)
+    except KeyboardInterrupt:
+        pass
+
+
 if __name__ == "__main__":
     main()