Add separate hooks for major state objects, add tests

hyperglass/external/bgptools.py

@@ -6,7 +6,6 @@
 
 # Standard Library
 import re
-import socket
 import asyncio
 from typing import Dict, List
 
@@ -87,54 +86,24 @@ async def run_whois(targets: List[str]) -> str:
     return response.decode()
 
 
-def run_whois_sync(targets: List[str]) -> str:
-    """Open raw socket to bgp.tools and execute query."""
-
-    # Construct bulk query
-    query = "\n".join(("begin", *targets, "end\n")).encode()
-
-    # Open the socket to bgp.tools
-    log.debug("Opening connection to bgp.tools")
-    sock = socket.socket()
-    sock.connect(("bgp.tools", 43))
-    sock.send(query)
-
-    # Read the response
-    response = b""
-    while True:
-        data = sock.recv(128)
-        if data:
-            response += data
-
-        else:
-            log.debug("Closing connection to bgp.tools")
-            sock.shutdown(1)
-            sock.close()
-            break
-
-    return response.decode()
-
-
 async def network_info(*targets: str) -> Dict[str, Dict[str, str]]:
     """Get ASN, Containing Prefix, and other info about an internet resource."""
 
     targets = [str(t) for t in targets]
-    (cache := use_state().redis)
+    cache = use_state("cache")
 
     # Set default data structure.
     data = {t: {k: "" for k in DEFAULT_KEYS} for t in targets}
 
     # Get all cached bgp.tools data.
-    cached = cache.hgetall(CACHE_KEY)
+    cached = cache.get_map(CACHE_KEY) or {}
 
     # Try to use cached data for each of the items in the list of
     # resources.
-    for t in targets:
-
-        if t in cached:
-            # Reassign the cached network info to the matching resource.
-            data[t] = cached[t]
-            log.debug("Using cached network info for {}", t)
+    for t in (t for t in targets if t in cached):
+        # Reassign the cached network info to the matching resource.
+        data[t] = cached[t]
+        log.debug("Using cached network info for {}", t)
 
     # Remove cached items from the resource list so they're not queried.
     targets = [t for t in targets if t not in cached]
@@ -149,7 +118,7 @@ async def network_info(*targets: str) -> Dict[str, Dict[str, str]]:
 
                 # Cache the response
                 for t in targets:
-                    cache.hset(CACHE_KEY, t, data[t])
+                    cache.set_map_item(CACHE_KEY, t, data[t])
                     log.debug("Cached network info for {}", t)
 
     except Exception as err:
@@ -160,42 +129,4 @@ async def network_info(*targets: str) -> Dict[str, Dict[str, str]]:
 
 def network_info_sync(*targets: str) -> Dict[str, Dict[str, str]]:
     """Get ASN, Containing Prefix, and other info about an internet resource."""
-
-    targets = [str(t) for t in targets]
-    (cache := use_state().redis)
-
-    # Set default data structure.
-    data = {t: {k: "" for k in DEFAULT_KEYS} for t in targets}
-
-    # Get all cached bgp.tools data.
-    cached = cache.hgetall(CACHE_KEY)
-
-    # Try to use cached data for each of the items in the list of
-    # resources.
-    for t in targets:
-
-        if t in cached:
-            # Reassign the cached network info to the matching resource.
-            data[t] = cached[t]
-            log.debug("Using cached network info for {}", t)
-
-    # Remove cached items from the resource list so they're not queried.
-    targets = [t for t in targets if t not in cached]
-
-    try:
-        if targets:
-            whoisdata = run_whois_sync(targets)
-
-            if whoisdata:
-                # If the response is not empty, parse it.
-                data.update(parse_whois(whoisdata, targets))
-
-                # Cache the response
-                for t in targets:
-                    cache.hset(CACHE_KEY, t, data[t])
-                    log.debug("Cached network info for {}", t)
-
-    except Exception as err:
-        log.error(str(err))
-
-    return data
+    return asyncio.run(network_info(*targets))