* Now possible to remove documents, new /delete endpoint.

* More unittests.

3 files changed, 113 insertions, 0 deletions

diff --git a/src/db.py b/src/db.py
index 7e83d96..143d0a5 100755
--- a/src/db.py
+++ b/src/db.py
@@ -121,3 +121,14 @@ class DictDB():
             data.append(doc)
 
         return data
+
+    def delete(self, key):
+        """
+        Delete a document based on its ID.
+        """
+        try:
+            self.couchdb.delete(key)
+        except couch.exceptions.NotFound:
+            return None
+
+        return key
diff --git a/src/main.py b/src/main.py
index 609433b..fef4050 100755
--- a/src/main.py
+++ b/src/main.py
@@ -24,6 +24,8 @@ app.add_middleware(
 )
 
 # TODO: X-Total-Count
+
+
 @app.middleware("http")
 async def mock_x_total_count_header(request: Request, call_next):
     response = await call_next(request)
@@ -153,6 +155,16 @@ async def add(data: Request, Authorize: AuthJWT = Depends()):
     return JSONResponse(content={"status": "success", "docs": key})
 
 
+@app.delete('/sc/v0/delete/{key}')
+async def delete(key):
+    if db.delete(key) is None:
+        return JSONResponse(content={"status": "error",
+                                     "message": "Document not found"},
+                            status_code=400)
+
+    return JSONResponse(content={"status": "success", "docs": {}})
+
+
 def main(standalone=False):
     if not standalone:
         return app
diff --git a/src/test/test_api.py b/src/test/test_api.py
index 0786363..cb3c02a 100644
--- a/src/test/test_api.py
+++ b/src/test/test_api.py
@@ -16,6 +16,8 @@ JWT_HEADER = {'Authorization': f'Bearer {JWT_TOKEN}'}
 
 
 def test_001():
+    print("*** Adding document.")
+
     doc_port = random.randint(1, 65536)
     doc_ip = str(ipaddress.IPv4Address(random.randint(1, 0xffffffff)))
     doc_asn = str(doc_ip) + '_' + str(doc_port)
@@ -65,13 +67,101 @@ def test_001():
 
 
 def test_002():
+    nr_documents = 100
+    starttime = time.time()
+
+    for i in range(nr_documents):
+        doc_port = random.randint(1, 65536)
+        doc_ip = str(ipaddress.IPv4Address(random.randint(1, 0xffffffff)))
+        doc_asn = str(doc_ip) + '_' + str(doc_port)
+
+        json_data = {
+            'ip': doc_ip,
+            'port': doc_port,
+            'whois_description': 'unittest',
+            'asn': doc_asn,
+            'asn_country_code': 'SE',
+            'ptr': 'unittest.example.com',
+            'abuse_mail': 'unittest@example.com',
+            'domain': 'sunet.se',
+            'timestamp_in_utc': '2021-06-21T14:06UTC',
+            'producer_unique_keys': {
+                'subject_cn': 'unittest',
+                'subject_o': 'unittest',
+                'full_name': 'unittest',
+                'end_of_general_support': False,
+                'cve_2021_21972': 'unittest',
+                'cve_2021_21974': 'unittest',
+                'cve_2021_21985': 'unittest'
+            }
+        }
+
+        response = client.post(
+            "/sc/v0/add", headers=JWT_HEADER, json=json_data)
+        assert(response.status_code == 200)
+        assert(response.json()['status'] == 'success')
+
+        response = client.get(
+            f"/sc/v0/get?port={doc_port}", headers=JWT_HEADER)
+        assert(response.status_code == 200)
+        assert(response.json()['status'] == 'success')
+        assert(len(response.json()['docs']) == 1)
+        assert(response.json()['docs'][0]['port'] == doc_port)
+
+        response = client.get(f"/sc/v0/get?asn={doc_asn}", headers=JWT_HEADER)
+        assert(response.status_code == 200)
+        assert(response.json()['status'] == 'success')
+        assert(len(response.json()['docs']) == 1)
+        assert(response.json()['docs'][0]['asn'] == doc_asn)
+
+        response = client.get(f"/sc/v0/get?ip={doc_ip}", headers=JWT_HEADER)
+        assert(response.status_code == 200)
+        assert(response.json()['status'] == 'success')
+        assert(len(response.json()['docs']) == 1)
+        assert(response.json()['docs'][0]['ip'] == doc_ip)
+
+    stop_time = str(time.time() - starttime)
+    print(f"*** Adding {nr_documents} documents took {stop_time} seconds.")
+
+
+def test_003():
     response = client.get("/sc/v0/get", headers=JWT_HEADER)
     assert(response.status_code == 200)
 
     for doc in response.json()['docs']:
         doc_id = doc['_id']
+
         response_doc = client.get(f"/sc/v0/get/{doc_id}", headers=JWT_HEADER)
         assert(response_doc.status_code == 200)
         assert(response_doc.json()['status'] == 'success')
         assert(type(response_doc.json()['docs']) == type(dict()))
         assert(response_doc.json()['docs']['domain'] == 'sunet.se')
+
+
+def test_004():
+    response = client.get("/sc/v0/get?limit=1000", headers=JWT_HEADER)
+    assert(response.status_code == 200)
+
+    nr_documents = len(response.json()['docs'])
+    starttime = time.time()
+
+    for doc in response.json()['docs']:
+        doc_id = doc['_id']
+        response_doc = client.delete(
+            f"/sc/v0/delete/{doc_id}", headers=JWT_HEADER)
+        assert(response_doc.status_code == 200)
+        assert(response_doc.json()['status'] == 'success')
+        response_doc = client.get(
+            f"/sc/v0/get/{doc_id}", headers=JWT_HEADER)
+        assert(response_doc.status_code == 200)
+        assert(response_doc.json()['status'] == 'success')
+        assert(response_doc.json()['docs'] == {})
+
+    stop_time = str(time.time() - starttime)
+    print(f"*** Removing {nr_documents} documents took {stop_time} seconds.")
+
+    print("*** Removing document with invalid ID.")
+    response = client.delete(
+        "/sc/v0/delete/nonexistent", headers=JWT_HEADER)
+    assert(response.status_code == 400)
+    assert(response.json()['status'] == 'error')