Amendments to auto-scan service and CronJob manifest (canada-ca#752)

* Amendments to auto-scan service and CronJob manifest to support scheduled scanning

* 8 hour interval -> 6 hour interval

Author: Ethanljf

platform/overlays/gke/knative/config/autoscan.yaml

@@ -2,9 +2,10 @@ apiVersion: batch/v1beta1
 kind: CronJob
 metadata:
   name: autoscan
-  namespace: tracker
+  namespace: scanners
 spec:
-  schedule: "0 2 * * *"
+  schedule: "*/360 * * * *"
+  concurrencyPolicy: Forbid
   jobTemplate:
     spec:
       template:

services/auto-scan/Dockerfile

@@ -8,13 +8,15 @@ COPY . ./
 RUN apt-get update
 RUN apt-get install -y --no-install-recommends apt-utils && \
 apt-get install -y --no-install-recommends python3 python3-pip && \
-apt-get install -y --no-install-recommends python3-setuptools python3-wheel
+apt-get install -y --no-install-recommends python3-setuptools python3-wheel && \
+apt-get install -y --no-install-recommends build-essential python3-dev
 
 # Install dependencies.
-RUN python3 -m pip install -r requirements.txt
+RUN pip3 install -r requirements.txt
+
+RUN chmod +x docker-entrypoint.sh
 
-RUN chmod +x autoscan.py
 RUN useradd -r -u 1001 autoscan
 USER autoscan
 
-ENTRYPOINT ["python3", "autoscan.py"]
+CMD exec ~/./docker-entrypoint.sh

services/auto-scan/autoscan.py

@@ -6,7 +6,7 @@
 import datetime
 import databases
 import sqlalchemy
-import bcrypt
+import traceback
 from sqlalchemy.sql import select
 from sqlalchemy.dialects.postgresql import ARRAY
 
@@ -19,7 +19,7 @@
 DB_HOST = os.getenv("DB_HOST")
 
 DATABASE_URI = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
-DISPATCHER_URL = "http://dispatcher.tracker.svc.cluster.local"
+QUEUE_URL = "http://scan-queue.scanners.svc.cluster.local"
 
 metadata = sqlalchemy.MetaData()
 
@@ -219,47 +219,35 @@
     sqlalchemy.Column("ref_links", sqlalchemy.String),
 )
 
-Classification = sqlalchemy.Table(
-    "classification",
-    metadata,
-    sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True),
-    sqlalchemy.Column("UNCLASSIFIED", sqlalchemy.String),
-)
-
 
-def Service(database=databases.Database(DATABASE_URI), client=requests):
+def Dispatch(database=databases.Database(DATABASE_URI), client=requests):
 
-    async def dispatch_web(domain, scan_id):
+    async def dispatch_https(domain, scan_id):
 
         payload = {
             "scan_id": scan_id,
             "domain": domain.get("domain"),
-            "selectors": domain.get("selectors"),
-            "user_init": False,
         }
-        headers = {
-            "Content-Type": "application/json",
-            "Data": str(payload),
-            "Scan-Type": "web",
+        client.post(QUEUE_URL+"/https", json=payload)
+
+    async def dispatch_ssl(domain, scan_id):
+
+        payload = {
+            "scan_id": scan_id,
+            "domain": domain.get("domain"),
         }
-        client.post(DISPATCHER_URL + "/receive", headers=headers)
+        client.post(QUEUE_URL+"/ssl", json=payload)
 
-    async def dispatch_mail(domain, scan_id):
+    async def dispatch_dns(domain, scan_id):
 
         payload = {
             "scan_id": scan_id,
             "domain": domain.get("domain"),
             "selectors": domain.get("selectors"),
-            "user_init": False,
-        }
-        headers = {
-            "Content-Type": "application/json",
-            "Data": str(payload),
-            "Scan-Type": "mail",
         }
-        client.post(DISPATCHER_URL + "/receive", headers=headers)
+        client.post(QUEUE_URL+"/dns", json=payload)
 
-    async def main():
+    async def scan():
         logging.info("Retrieving domains for scheduled scan...")
         try:
             await database.connect()
@@ -273,8 +261,14 @@ async def main():
             system = await database.fetch_one(system_user_query)
 
             scan_time = datetime.datetime.utcnow()
+            count = 0
 
             for domain in domains:
+                count = count + 1
+                dispatched.append(domain.get('domain'))
+                logging.info(f"Dispatching scan number {count} of {len(domains)}")
+                logging.info(f"Requesting scan for {domain.get('domain')}")
+
                 web_insert = Web_scans.insert().values(
                     domain_id=domain.get("id"),
                     scan_date=scan_time,
@@ -297,25 +291,23 @@ async def main():
                 web_scan_query = select([Web_scans]).order_by(Web_scans.c.id.desc())
                 web_scan = await database.fetch_one(web_scan_query)
 
-                dispatched.append(dispatch_web(domain, web_scan.get("id")))
-                dispatched.append(dispatch_mail(domain, mail_scan.get("id")))
-
-            for domain in dispatched:
-                await domain
+                await dispatch_https(domain, web_scan.get("id"))
+                await dispatch_ssl(domain, web_scan.get("id"))
+                await dispatch_dns(domain, mail_scan.get("id"))
 
             await database.disconnect()
-            return {"Dispatched": [domain.get('domain') for domain in domains]}
+            return [domain for domain in dispatched]
 
         except Exception as e:
             try:
                 await database.disconnect()
             except:
                 pass
-            logging.error(e)
-            return {"Dispatched": []}
-
-    return asyncio.run(main())
+            logging.error(f"An unexpected error occurred while initiating scheduled scan: {str(e)}\n\nFull traceback: {traceback.format_exc()}")
+            return [domain for domain in dispatched]
 
+    return asyncio.run(scan())
 
 if __name__ == "__main__":
-    logging.info(Service())
+    dispatched_domains = Dispatch()
+    logging.info(f"Dispatched scans for: {str(dispatched_domains)}")

services/auto-scan/docker-entrypoint.sh

@@ -0,0 +1,4 @@
+#!/bin/sh
+
+sleep 20
+python3 autoscan.py

services/auto-scan/requirements.txt

@@ -1,8 +1,7 @@
 asyncio
 pytest-asyncio
 requests>=2.18.4
-graphene
 SQLAlchemy
+psycopg2-binary
 databases[postgresql]
 pretend
-bcrypt

services/auto-scan/tests/test_autoscan.py

@@ -4,7 +4,7 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 from pretend import stub
-from autoscan import Service, Domains, Users
+from autoscan import Dispatch, Domains, Users
 
 TEST_DATABASE_URI = "postgresql://track_dmarc:postgres@testdb/track_dmarc"
 engine = create_engine(TEST_DATABASE_URI)
@@ -27,8 +27,8 @@ def test_retrieval():
             test_session.execute(Domains.insert().values(domain=domain["domain"]))
             test_session.execute(Domains.insert().values(domain=domain["domain"]))
 
-    client_stub = stub(post=lambda url, headers: None)
+    client_stub = stub(post=lambda url, json: None)
 
-    response = Service(database=databases.Database(TEST_DATABASE_URI), client=client_stub)
+    dispatched = Dispatch(database=databases.Database(TEST_DATABASE_URI), client=client_stub)
 
-    assert all(domain["domain"] in response["Dispatched"] for domain in input_domains)
+    assert all(domain["domain"] in dispatched for domain in input_domains)