chg: [library] moves things around

private_search_set/cli.py

@@ -1,50 +1,12 @@
 import click
-import json
-import os
-import sys
-import pdb
 from private_search_set.main import PrivateSearchSet
-from flor import BloomFilter
-
-def json_specs(json_file):
-    with open(json_file) as file:
-        data = convert_hyphen_in_json(json_file)
-        pss = PrivateSearchSet(**data)  # Create an instance of the PrivateSearchSet class
-    if set(data.keys()) == set(pss.__dict__.keys()):
-        click.echo("Fields in JSON file correspond to the private search set class.")
-        PrivateSearchSet.print_private_search_set(pss)
-        return pss
-    else:
-        click.echo("Fields in JSON file do not correspond to the private search set class.")
-        exit(1)
-
-def convert_hyphen_in_json(json_file):
-    """Convert hyphens in JSON file to underscores."""
-    with open(json_file) as file:
-        data = json.load(file)
-        data = {k.replace('-', '_'): v for k, v in data.items()}
-        return data
 
 @click.pass_context
 def ingest_stdin(ctx):
     pss = ctx.obj
     click.echo("Ingesting stdin to PSS file.")
-    # Create the folder if it does not exist
-    if not os.path.exists(ctx.params["pss_home"]):
-        os.makedirs(ctx.params["pss_home"])
-    if pss.bloomfilter['format'] == 'dcso-v1':
-        bf = BloomFilter(n=pss.bloomfilter['capacity'], p=pss.bloomfilter['fp-probability'])
-        # Read from stdin  
-        for line in sys.stdin.buffer.read().splitlines():  
-            click.echo(line)
-            bf.add(line)
-        file_path = os.path.join(ctx.params["pss_home"], 'private-search-set.bloom')
-        with open(file_path, 'wb') as f:
-            bf.write(f)
-    else:
-        print("Bloomfilter format not supported.")
-        exit(1) 
-
+    pss.ingest_stdin()
+    pss.write_to_files(ctx.params["pss_home"])
 
 
 @click.command()
@@ -54,7 +16,7 @@ def ingest_stdin(ctx):
 @click.option('--debug/--no-debug', default=False)
 @click.pass_context
 def cli(ctx, json_file, pss_home, ingest, debug):
-    ctx.obj = json_specs(json_file)
+    ctx.obj = PrivateSearchSet.load_from_json_specs(json_file)
     if ingest:
         ingest_stdin()
     pass

private_search_set/main.py

@@ -1,3 +1,8 @@
+import json 
+import os
+from flor import BloomFilter
+import sys
+
 class PrivateSearchSet:
     def __init__(self, algorithm, bloomfilter, canonicalization_format, description, generated_timestamp, keyid, misp_attribute_types, version):
         self.algorithm = algorithm
@@ -19,6 +24,34 @@ def print_private_search_set(private_search_set):
         print("MISP Attribute Types:", private_search_set.misp_attribute_types)
         print("Version:", private_search_set.version)
 
+    def load_from_json_specs(json_file):
+        with open(json_file) as file:
+            json_data = json.load(file)
+            data = {k.replace('-', '_'): v for k, v in json_data.items()}
+            pss = PrivateSearchSet(**data)  # Create an instance of the PrivateSearchSet class
+        if set(data.keys()) == set(pss.__dict__.keys()):
+            PrivateSearchSet.print_private_search_set(pss)
+            return pss
+        else:
+            raise ValueError("JSON file does not match the expected format.")
+
+    def ingest_stdin(self):
+        if self.bloomfilter['format'] == 'dcso-v1':
+            self.bf = BloomFilter(n=self.bloomfilter['capacity'], p=self.bloomfilter['fp-probability'])
+            # Read bytes from stdin  
+            for line in sys.stdin.buffer.read().splitlines():  
+                self.bf.add(line)
+        else:
+            raise ValueError("Bloomfilter format not supported.")
+
+    def write_to_files(self, pss_home):
+        if not os.path.exists(pss_home):
+            os.makedirs(pss_home)
+        file_path = os.path.join(pss_home, 'private-search-set.bloom')
+        with open(file_path, 'wb') as f:
+            self.bf.write(f)
+
+
 # Example usage:
 # json_data = {
 #     "algorithm": "Blake2",