# Buggy code
def get_web_screenshots(target, scan_id, proctimeout):
    scan_dir = utils.get_scan_dir(scan_id)
    xml_file = os.path.join(scan_dir, f"nmap.{scan_id}.xml")
    outFiles = os.path.join(scan_dir, f"aquatone.{scan_id}")
    output = []
    logger.info(f"Attempting to take screenshots for {target}")

    aquatoneArgs = ["aquatone", "-nmap", "-scan-timeout", "2500", "-out", outFiles]
    with open(xml_file, "r") as f:
        process = subprocess.Popen(aquatoneArgs, stdin=f, stdout=subprocess.DEVNULL)  # nosec

    try:
        process.communicate(timeout=proctimeout)
        if process.returncode == 0:
            time.sleep(
                0.5
            )  # a small sleep to make sure all file handles are closed so that the agent can read them
    except subprocess.TimeoutExpired:
        logger.warning(f"TIMEOUT: Killing aquatone against {target}")
        process.kill()

    session_path = os.path.join(outFiles, "aquatone_session.json")
    if not os.path.isfile(session_path):
        return output

    with open(session_path) as f:
        session = json.load(f)

    if session["stats"]["screenshotSuccessful"] > 0:
        logger.info(
            f"{target} - Success: {session['stats']['screenshotSuccessful']}, Fail: {session['stats']['screenshotFailed']}"
        )

        for k, page in session["pages"].items():
            fqScreenshotPath = os.path.join(outFiles, page["screenshotPath"])
            if page["hasScreenshot"] and os.path.isfile(fqScreenshotPath):
                urlp = urlparse(page["url"])
                if not urlp.port and urlp.scheme == "http":
                    port = 80
                elif not urlp.port and urlp.scheme == "https":
                    port = 443
                else:
                    port = urlp.port
                logger.info(
                    f"{urlp.scheme.upper()} screenshot acquired for {page['hostname']} on port {port}"
                )
                output.append(
                    {
                        "host": page["hostname"],
                        "port": port,
                        "service": urlp.scheme.upper(),
                        "data": base64_image(fqScreenshotPath),
                    }
                )
    return output


# Traceback (PIL.UnidentifiedImageError)
[2020-07-25 00:22:59,784] ERROR in app: Exception on /api/submit [POST]
Traceback (most recent call last):
File "/.venv/lib/python3.8/site-packages/flask/app.py", line 2447, in wsgi_app
response = self.full_dispatch_request()
File "/.venv/lib/python3.8/site-packages/flask/app.py", line 1952, in full_dispatch_request
rv = self.handle_user_exception(e)
File "/.venv/lib/python3.8/site-packages/flask/app.py", line 1821, in handle_user_exception
reraise(exc_type, exc_value, tb)
File "/.venv/lib/python3.8/site-packages/flask/_compat.py", line 39, in reraise
raise value
File "/.venv/lib/python3.8/site-packages/flask/app.py", line 1950, in full_dispatch_request
rv = self.dispatch_request()
File "/.venv/lib/python3.8/site-packages/flask/app.py", line 1936, in dispatch_request
return self.view_functions[rule.endpoint](**req.view_args)
File "/opt/natlas/natlas-server/app/auth/wrappers.py", line 49, in decorated_function
return f(*args, **kwargs)
File "/opt/natlas/natlas-server/app/api/routes.py", line 187, in submit
newhost["screenshots"], newhost["num_screenshots"] = process_screenshots(
File "/opt/natlas/natlas-server/app/api/processing/screenshot.py", line 52, in process_screenshots
item["thumb_hash"] = create_thumbnail(fname, file_ext)
File "/opt/natlas/natlas-server/app/api/processing/screenshot.py", line 10, in create_thumbnail
thumb = Image.open(fname)
File "/.venv/lib/python3.8/site-packages/PIL/Image.py", line 2895, in open
raise UnidentifiedImageError(
PIL.UnidentifiedImageError: cannot identify image file '/data/media/original/e3/b0/e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855.png'

# Fixed code
def get_web_screenshots(target, scan_id, proctimeout):
    scan_dir = utils.get_scan_dir(scan_id)
    xml_file = os.path.join(scan_dir, f"nmap.{scan_id}.xml")
    output_dir = os.path.join(scan_dir, f"aquatone.{scan_id}")
    logger.info(f"Attempting to take screenshots for {target}")

    aquatoneArgs = ["aquatone", "-nmap", "-scan-timeout", "2500", "-out", output_dir]
    with open(xml_file, "r") as f:
        process = subprocess.Popen(aquatoneArgs, stdin=f, stdout=subprocess.DEVNULL)  # nosec

    try:
        process.communicate(timeout=proctimeout)
        if process.returncode == 0:
            time.sleep(
                0.5
            )  # a small sleep to make sure all file handles are closed so that the agent can read them
    except subprocess.TimeoutExpired:
        logger.warning(f"TIMEOUT: Killing aquatone against {target}")
        process.kill()

    return parse_aquatone_session(output_dir)