from scapy.all import TCP

class Extracteur:
    def __init__(self, fichier_pcap):
        pcap = rdpcap(fichier_pcap)
        self.sessions = pcap.sessions()
        self.responses = []

    def trouver_reponses(self):
        for session in self.sessions:
            payload = b""
            for paquet in self.sessions[session]:
                try:
                    if paquet[TCP].dport == 80 or paquet[TCP].sport == 80:
                        payload += bytes(paquet[TCP].payload)
                except IndexError:
                    pass

            if payload:
                header = extract_header(payload)
                if header is None:
                    continue
                self.responses.append(Response(header=header, payload=payload))

    def ecrire_images(self):
        os.makedirs(DOSSIER_SORTIE, exist_ok=True)
        compteur = 0
        for response in self.responses:
            contenu, content_type = extract_content(response)
            if contenu and content_type:
                nom = os.path.join(DOSSIER_SORTIE, f"image_{compteur}.{content_type}")
                with open(nom, "wb") as f:
                    f.write(contenu)
                print(f"  {nom}")
                compteur += 1
        print(f"\n{compteur} images extraites")