async def test_scan_db():
# TO DO: Add tests for all ScanDatabase methods
os.mkdir(report_folder)
try:
await ScanDatabase.create_db_and_schema(report_folder)
async with ScanDatabase(report_folder=report_folder) as db:
fake_data = {
"ip": "127.0.0.1",
"hostname": "test",
"url": "https://127.0.0.1:443/",
"screenshot": f"{report_folder}/test.png",
"port": 443,
"scheme": "https",
"title": "Test page",
"server": "Test server",
"headers": {"Wat": "Test"},
"body": "<html><body> Test </body></html>",
}
await db.add_host_and_service(**fake_data)
assert await db.get_host_count() == 1
assert await db.get_service_count() == 1
assert len(await db.get_hosts()) == 1
assert len(await db.get_services()) == 1
finally:
shutil.rmtree(report_folder, ignore_errors=True)
async def _print_hosts(self, hosts, table_title=None):
table_data = [["Id", "IP", "Hostname", "Discovered Services", "Matched Signature(s)"]]
async with ScanDatabase(connection=self.db) as db:
for entry in hosts:
host_id, hostname, ip = entry
service_count = await db.get_service_count_on_host(host_id)
matched_sigs = map(
lambda x: x[0].split(','),
filter(
lambda x: x[0] is not None,
await db.get_matched_sigs_on_host(host_id)
)
)
table_data.append([
host_id,
ip,
hostname,
service_count,
','.join(set(sig_name for result in matched_sigs for sig_name in result))
])
table = AsciiTable(table_data)
table.inner_row_border = True
table.title = table_title
print(table.table)
async def scan(self):
"""
Peform a signature scan on all discovered servers
"""
self.signatures.load()
log.debug("Starting signature scan...")
start_time = time()
async with ScanDatabase(connection=self.db) as db:
tasks = [
self.signatures.find_match(service)
for service in await db.get_services()
]
matches = list(
filter(lambda x: len(x[0]) > 0, await asyncio.gather(*tasks)))
for match in matches:
await db.add_matched_sigs_to_service(
match[1][0], ",".join([sig["name"] for sig in match[0]]))
completed_time = strftime("%Mm%Ss", gmtime(time() - start_time))
log.debug(
f"Signature scan completed, identified {len(matches)} service(s) in {completed_time}"
)
async def worker(browser, queue):
while True:
url = await queue.get()
page = await browser.newPage()
page.setDefaultNavigationTimeout(
args.timeout *
100) # setDefaultNavigationTimeout() accepts milliseconds
#page.on('request', lambda req: asyncio.create_task(on_request(req)))
#page.on('requestfinished', lambda req: asyncio.create_task(on_requestfinished(req)))
#page.on('response', lambda resp: asyncio.create_task(on_response(resp)))
try:
r = await asyncio.wait_for(screenshot(url, page),
timeout=args.timeout)
logging.info(r)
async with ScanDatabase(report_folder) as db:
await db.add_host_and_service(**r)
logging.info(f"Took screenshot of {url}")
except asyncio.TimeoutError:
logging.info(f"Task for url {url} timed out")
except Exception as e:
#if not any(err in str(e) for err in ['ERR_ADDRESS_UNREACHABLE', 'ERR_CONNECTION_REFUSED', 'ERR_CONNECTION_TIMED_OUT']):
logging.error(f"Error taking screenshot: {e}")
finally:
stats.execs += 1
await page.close()
queue.task_done()
async def show(self, *args, **kwargs):
"""
Preview screenshot in Terminal
"""
async with ScanDatabase(connection=self.db) as db:
entry = await db.get_service_by_id(int(*args[0]))
_,_,screenshot_path,_,_,_,_,_,_ = entry
imgcat(open(screenshot_path))
async def open(self, *args, **kwargs):
"""
Open screenshot in browser/previewer
"""
async with ScanDatabase(connection=self.db) as db:
entry = await db.get_service_by_id(int(*args[0]))
_,_,screenshot_path,_,_,_,_,_,_ = entry
screenshot_path = str(pathlib.Path(screenshot_path).absolute())
webbrowser.open(screenshot_path.replace("/", "file:////", 1))
async def servers(self, args):
"""
Show discovered servers
"""
async with ScanDatabase(connection=self.db) as db:
if len(args):
query_results = await db.search_services(args[0])
else:
query_results = await db.get_services()
await self._print_services(query_results)
async def hosts(self, *args, **kwargs):
"""
Show hosts
"""
async with ScanDatabase(connection=self.db) as db:
table_data = [["Id", "IP", "Hostname"]]
for entry in await db.get_hosts():
host_id, hostname, ip = entry
table_data.append([host_id, ip, hostname])
table = AsciiTable(table_data)
table.inner_row_border = True
print(table.table)
async def servers(self, *args, **kwargs):
"""
Show discovered servers
"""
async with ScanDatabase(connection=self.db) as db:
table_data = [["Id", "URL", "Title", "Server"]]
for entry in await db.get_services():
service_id, url, title, headers = entry
table_data.append(
[service_id, url, title,
json.loads(headers)['server']])
table = AsciiTable(table_data)
table.inner_row_border = True
print(table.table)
async def open(self, args):
"""
Open a screenshot in your default browser/previewer
"""
try:
server_id = int(args[0])
except IndexError:
print("No server id given")
except ValueError:
print("Invalid server id")
else:
async with ScanDatabase(connection=self.db) as db:
entry = await db.get_service_by_id(server_id)
screenshot_path = self.scan_folder_path / entry[2]
webbrowser.open(screenshot_path.absolute().as_uri())
async def open(self, args):
"""
Open a screenshot in your default browser/previewer
"""
try:
server_id = int(args[0])
except IndexError:
print("No server id given")
except ValueError:
print("Invalid server id")
else:
async with ScanDatabase(connection=self.db) as db:
entry = await db.get_service_by_id(server_id)
screenshot_path = str(db_path.parent.joinpath(entry[2]).absolute())
webbrowser.open(screenshot_path.replace("/", "file:////", 1))
async def show(self, args):
"""
Preview a screenshot in the terminal
"""
try:
server_id = int(args[0])
except IndexError:
print("No server id given")
except ValueError:
print("Invalid server id")
else:
async with ScanDatabase(connection=self.db) as db:
entry = await db.get_service_by_id(server_id)
imgcat(
open(db_path.parent.joinpath(entry[2]).absolute())
)
async def show(self, args):
"""
Preview a screenshot in the terminal
"""
try:
server_id = int(args[0])
except IndexError:
print("No server id given")
except ValueError:
print("Invalid server id")
else:
async with ScanDatabase(connection=self.db) as db:
entry = await db.get_service_by_id(server_id)
with open((self.scan_folder_path / entry[2]).absolute(),
"rb") as image:
imgcat(image)
async def generate_csv_report(scan_folder, db):
result_limit = 100
csv_path = scan_folder / "witnessme_report.csv"
with open(csv_path, "w", newline="") as csvfile:
writer = csv.writer(csvfile)
writer.writerow([
"url",
"ip",
"hostname",
"port",
"scheme",
"title",
"server",
"screenshot",
"matched_sigs",
])
log.info("Generating CSV report, please wait...")
async with ScanDatabase(connection=db) as db:
offset = 0
while True:
services = await db.get_services_with_host(limit=result_limit,
offset=offset)
if not services:
break
for service in services:
writer.writerow([
service[1],
service[12],
service[11],
service[3],
service[4],
service[5],
service[6],
service[2],
service[9],
])
offset += result_limit
log.info("Done")
async def hosts(self, args):
"""
Show hosts
"""
async with ScanDatabase(connection=self.db) as db:
try:
filter_term = args[0]
except IndexError:
hosts = await db.get_hosts()
await self._print_hosts(hosts)
else:
try:
host = await db.get_host_by_id(int(filter_term))
if not host: raise ValueError(f"No host found with id: {filter_term}")
except ValueError:
query_results = await db.search_hosts(filter_term)
await self._print_hosts(query_results)
else:
await self._print_hosts([host])
services = await db.get_services_on_host(host[0])
await self._print_services(services)
async def generate_html_report(scan_folder, db):
results_per_page = 100
template_path = pkg_resources.resource_filename(__name__,
"templates/template.html")
log.info("Generating HTML report, please wait...")
async with ScanDatabase(connection=db) as db:
service_count = await db.get_service_count()
total_pages = (service_count // results_per_page if service_count %
results_per_page == 0 else
(service_count // results_per_page) + 1)
current_page = 1
offset = 0
while True:
services = await db.get_services_with_host(limit=results_per_page,
offset=offset)
if not services:
break
report_file = scan_folder / f"report_page_{current_page}.html"
if current_page == 1:
report_file = scan_folder / "witnessme_report.html"
with open(report_file, "w") as report:
with open(template_path) as file_:
template = Template(file_.read())
report.write(
template.render(
name="WitnessMe Report",
current_page=current_page,
total_pages=total_pages,
services=services,
))
current_page += 1
offset += results_per_page
log.info("Done")
async def add_to_database(self, browser, results):
log.info(f"Took screenshot of {results['url']}")
async with ScanDatabase(self.report_folder) as db:
await db.add_host_and_service(**results)