From fd91b827e977428e44d9404c2de084b14047e5b2 Mon Sep 17 00:00:00 2001
From: Jonathan Lundy <jonathan.lundy@rakuten.com>
Date: Tue, 23 Jun 2020 10:38:35 -0600
Subject: [PATCH] add scan-dir and scan-file to rspl

---
 utils/registry/dom/filedom.py        |  2 ++
 utils/registry/main.py               |  4 +++-
 utils/registry/rpsl_scan/__init__.py | 36 ++++++++++++++++++++++------
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/utils/registry/dom/filedom.py b/utils/registry/dom/filedom.py
index b775433c6..9f3186e29 100644
--- a/utils/registry/dom/filedom.py
+++ b/utils/registry/dom/filedom.py
@@ -146,6 +146,8 @@ class FileDOM:
         self.keys = keys
         self.multi = multi
         self.mntner = mntner
+        if self.src is None:
+            self.src = f"{self.schema}/{self.name}"
         self.valid = True
 
     @property
diff --git a/utils/registry/main.py b/utils/registry/main.py
index 1b193f767..863beb9a9 100644
--- a/utils/registry/main.py
+++ b/utils/registry/main.py
@@ -70,6 +70,8 @@ def run() -> int:
     working_dir = os.getcwd()
     working_dir = os.environ.get("WORKING_DIR", working_dir)
     prog_dir = os.path.dirname(os.path.realpath(__file__))
+    rpsl_dir = os.environ.get("RPSL_DIR", working_dir)
+    rpsl_dir = find_rpsl(rpsl_dir)
 
     cmd, args = shift(shift(sys.argv)[1])
 
@@ -90,7 +92,7 @@ def run() -> int:
     return pkg.run(args, {
         "WORKING_DIR": working_dir,
         "BIN_DIR": prog_dir,
-        "RPSL_DIR": find_rpsl(working_dir),
+        "RPSL_DIR": rpsl_dir,
         })
 
 
diff --git a/utils/registry/rpsl_scan/__init__.py b/utils/registry/rpsl_scan/__init__.py
index dd9ec3a15..ff668b154 100644
--- a/utils/registry/rpsl_scan/__init__.py
+++ b/utils/registry/rpsl_scan/__init__.py
@@ -5,12 +5,18 @@
 
 import os
 import sys
+import argparse
 from typing import List, Dict
 
 from dom.filedom import FileDOM
 from dom.schema import SchemaDOM
 from dom.transact import TransactDOM
 
+parser = argparse.ArgumentParser()
+parser.add_argument("--add-index", action='store_true')
+parser.add_argument("--scan-dir", type=str, default=None)
+parser.add_argument("--scan-file", type=str, default=None)
+
 
 def index_files(path: str):
     """generate list of dom files"""
@@ -27,13 +33,20 @@ def index_files(path: str):
 
 def run(args: List[str], env: Dict[str, str]) -> int:
     """run scan script"""
+    opts = parser.parse_args(args)
 
     path = env.get("RPSL_DIR")
     if path is None:
-        print("RPSL index has not been generated.", file=sys.stderr)
+        print("RPSL directory not found. do `rpsl init` or set RPSL_DIR",
+              file=sys.stderr)
         return 1
 
     index_file = os.path.join(path, ".rpsl/index")
+    schema_file = os.path.join(path, ".rpsl/schema")
+
+    if not os.path.exists(index_file) or not os.path.exists(schema_file):
+        print("RPSL index files not found. do `rpsl index`?")
+        return 1
 
     lookups = {}  # type: Dict[str, FileDOM]
     schemas = {}  # type: Dict[str, SchemaDOM]
@@ -45,18 +58,27 @@ def run(args: List[str], env: Dict[str, str]) -> int:
             lookups[(sp[0], sp[1])] = (sp[2], "")
         print("done.", file=sys.stderr, flush=True)
 
-    schema_file = os.path.join(path, ".rpsl/schema")
     schema_set = TransactDOM.from_file(schema_file)
 
     for schema in schema_set.schemas:
         schemas[schema.ref] = schema
 
-    files = index_files(path)
-    # for dom in files:
-    #     key, value = dom.index
-    #     lookups[key] = value
+    def file_gen():
+        if opts.scan_dir is not None:
+            path = os.path.join(env.get("WORKING_DIR"), opts.scan_dir)
+        elif opts.scan_file is not None:
+            path = os.path.join(env.get("WORKING_DIR"), opts.scan_file)
+            return TransactDOM.from_file(path).files
 
-    for dom in files:
+        return index_files(path)
+
+    if opts.add_index:
+        print("Add scanned items to lookup index...", file=sys.stderr)
+        for dom in file_gen():
+            key, value = dom.index
+            lookups[key] = value
+
+    for dom in file_gen():
         s = schemas.get(dom.rel)
         if s is None:
             print(f"{dom.src} schema not found for {dom.rel}")