75 lines
5.1 KiB
Python
75 lines
5.1 KiB
Python
import json
|
|
import os
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from http.client import HTTPConnection
|
|
from tempfile import TemporaryDirectory
|
|
from unittest.mock import patch
|
|
|
|
from app.main import create_server, hash_password
|
|
|
|
|
|
class ScopedDiscoveryApiTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = TemporaryDirectory()
|
|
os.environ['BOOTSTRAP_ADMIN_EMAIL'] = 'discover-owner@example.test'
|
|
os.environ['BOOTSTRAP_ADMIN_PASSWORD'] = 'password'
|
|
self.server = create_server('127.0.0.1', 0, self.tmp.name + '/db.sqlite')
|
|
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True); self.thread.start()
|
|
self.conn = HTTPConnection('127.0.0.1', self.server.server_port, timeout=4); self.cookie = None
|
|
self.request('POST', '/api/v1/auth/login', {'email': 'discover-owner@example.test', 'password': 'password'})
|
|
|
|
def tearDown(self):
|
|
self.server.shutdown(); self.server.server_close(); self.thread.join(2); self.tmp.cleanup()
|
|
|
|
def request(self, method, path, payload=None):
|
|
body = json.dumps(payload).encode() if payload is not None else None
|
|
headers = {'Content-Type': 'application/json'} if body else {}
|
|
if self.cookie: headers['Cookie'] = self.cookie
|
|
self.conn.request(method, path, body, headers); response = self.conn.getresponse()
|
|
cookie = response.getheader('Set-Cookie')
|
|
if cookie: self.cookie = cookie.split(';', 1)[0]
|
|
return response.status, json.loads(response.read() or b'{}')
|
|
|
|
def test_direct_criteria_job_crawls_allowlisted_site_and_persists_evidence(self):
|
|
pages = {
|
|
'https://directory.test/': {'status': 200, 'final_url': 'https://directory.test/', 'content_type': 'text/html', 'body': b'<h1>Acme Solar</h1><a href="https://acme.test/">Acme</a>'},
|
|
'https://acme.test/': {'status': 200, 'final_url': 'https://acme.test/', 'content_type': 'text/html', 'body': b'<title>Acme Solar</title><h1>Acme Solar</h1><p>Solar installers</p><a href="/contact">Contact</a>'},
|
|
'https://acme.test/contact': {'status': 200, 'final_url': 'https://acme.test/contact', 'content_type': 'text/html', 'body': b'<h1>Contact Acme</h1><a href="mailto:hello@acme.test">Email</a><p>+27 11 555 0100</p>'},
|
|
}
|
|
def fetch(url, **_):
|
|
value = pages[url]; return dict(value, redirect_chain=[], elapsed_ms=1, tls=url.startswith('https://'), certificate_status='valid')
|
|
with patch('app.discovery._fetch', side_effect=fetch), patch('app.discovery.validate_url', side_effect=lambda url, **_: url), patch('app.main.validate_url', side_effect=lambda url, **_: url):
|
|
status, created = self.request('POST', '/api/v1/discovery', {'criteria': {'keywords': ['solar'], 'location': 'Cape Town'}, 'seed_urls': ['https://directory.test/'], 'max_pages': 5, 'idempotency_key': 'scope-1'})
|
|
self.assertEqual(status, 202); self.assertEqual(created['type'], 'scoped_discovery')
|
|
for _ in range(50):
|
|
_, job = self.request('GET', '/api/v1/jobs/' + str(created['id']))
|
|
if job['status'] in ('succeeded', 'failed'): break
|
|
time.sleep(.02)
|
|
self.assertEqual(job['status'], 'succeeded')
|
|
status, runs = self.request('GET', '/api/v1/discovery-runs')
|
|
self.assertEqual(status, 200); self.assertEqual(runs['items'][0]['criteria']['location'], 'Cape Town')
|
|
self.assertEqual(runs['items'][0]['result_count'], 1)
|
|
status, businesses = self.request('GET', '/api/v1/businesses')
|
|
self.assertEqual(status, 200); self.assertEqual(businesses['items'][0]['website_domain'], 'acme.test')
|
|
detail = self.request('GET', '/api/v1/businesses/' + str(businesses['items'][0]['id']))[1]
|
|
self.assertTrue(any(x['url'] == 'https://acme.test/contact' for x in detail['evidence']))
|
|
self.assertTrue(any(x['source_url'] == 'https://acme.test/contact' and x['provenance'] == 'mailto' for x in detail['contact_extractions']))
|
|
self.assertFalse(detail.get('outreach_enabled', False))
|
|
|
|
def test_requires_bounded_seed_allowlist_and_rejects_ssrf(self):
|
|
self.assertEqual(self.request('POST', '/api/v1/discovery', {'criteria': {'keywords': ['x']}})[0], 400)
|
|
status, body = self.request('POST', '/api/v1/discovery', {'criteria': {}, 'seed_urls': ['http://127.0.0.1/'], 'idempotency_key': 'bad'})
|
|
self.assertEqual(status, 400); self.assertEqual(body['error'], 'unsafe_seed_url')
|
|
|
|
def test_results_are_tenant_isolated(self):
|
|
ph, salt = hash_password('other-password')
|
|
db = sqlite3.connect(self.server.db_path); db.execute("INSERT INTO organizations VALUES ('other-tenant','Other',CURRENT_TIMESTAMP)"); db.execute("INSERT INTO users (organization_id,email,password_hash,password_salt,role) VALUES (?,?,?,?,?)", ('other-tenant','other@example.test',ph,salt,'owner')); db.commit(); db.close()
|
|
self.cookie = None; self.request('POST', '/api/v1/auth/login', {'email': 'other@example.test', 'password': 'other-password'})
|
|
self.assertEqual(self.request('GET', '/api/v1/discovery-runs')[1]['items'], [])
|
|
|
|
|
|
if __name__ == '__main__': unittest.main()
|