import csv
import time
import os
from gvm.connections import TLSConnection
from gvm.protocols.gmp import Gmp
from gvm.transforms import EtreeTransform

# Set up data directory
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.abspath(os.path.join(BASE_DIR, "..", "data"))
os.makedirs(DATA_DIR, exist_ok=True)

TASK_ID_CSV = os.path.join(DATA_DIR, "task_id.csv")
REPORT_ID_CSV = os.path.join(DATA_DIR, "report_id.csv")

OPENVAS_HOST = "localhost"
OPENVAS_PORT = 9390
USERNAME = "admin"
PASSWORD = "admin"

def wait_for_task_and_get_report(gmp, task_id):
    while True:
        response = gmp.get_task(task_id=task_id)
        task = response.find("task")

        if task is None:
            print(f"Task with ID {task_id} not found. Skipping.")
            return None

        status = task.find("status").text
        progress = task.find("progress").text
        print(f"Task {task_id} - Status: {status}, Progress: {progress}%")


        if status == "Done":
            report_elem = task.find("last_report/report")
            if report_elem is not None:
                report_id = report_elem.get("id")
                print(f"[✓] Task {task_id} finished. Report ID: {report_id}")
                return report_id
            else:
                print(f"[✗] Task {task_id} is done, but no report found.")
                return None
        time.sleep(10)

def main():
    connection = TLSConnection(hostname=OPENVAS_HOST, port=OPENVAS_PORT)
    with Gmp(connection=connection, transform=EtreeTransform()) as gmp:
        gmp.authenticate(username=USERNAME, password=PASSWORD)
        print("Authenticated with OpenVAS")

        try:
            with open(TASK_ID_CSV, newline="") as csvfile:
                reader = csv.DictReader(csvfile)
                task_ids = [row["Task ID"] for row in reader]

            with open(REPORT_ID_CSV, "w", newline="") as outfile:
                writer = csv.writer(outfile)
                writer.writerow(["Task ID", "Report ID"])  # Header

                for task_id in task_ids:
                    print(f"Waiting for report from Task ID: {task_id}")
                    report_id = wait_for_task_and_get_report(gmp, task_id)
                    if report_id:
                        writer.writerow([task_id, report_id])
        except FileNotFoundError:
            print(f"File not found: {TASK_ID_CSV}. Please run taskmaker.py first.")
        except KeyError:
            print("Error: CSV must contain a 'Task ID' column.")

if __name__ == "__main__":
    main()