from elasticsearch_dsl import connections, Search, Q

from parsedmarc.elastic import _AggregateReportDoc, _ForensicReportDoc
from parsedmarc.utils import get_ip_address_country


# Replace with your Elasticsearch URLs
connections.create_connection(hosts=["127.0.0.1:9200"])

search = Search(index="dmarc_aggregate*")

query = ~Q(dict(exists=dict(field="source_country")))

search.query = query

count = search.count()
search = search[0:count]
results = search.execute()

for result in results:
    doc = _AggregateReportDoc.get(id=result.meta.id, index=result.meta.index)
    source_ip_address = str(result.source_ip_address)
    source_country = get_ip_address_country(source_ip_address)
    if source_country:
        doc.source_country = source_country
        doc.save()


search = Search(index="dmarc_forensic*")

query = ~Q(dict(exists=dict(field="source_country")))

search.query = query

count = search.count()
search = search[0:count]
results = search.execute()

for result in results:
    doc = _ForensicReportDoc.get(id=result.meta.id, index=result.meta.index)
    source_ip_address = str(result.source_ip_address)
    source_country = get_ip_address_country(source_ip_address)
    if source_country:
        doc.source_country = source_country
        doc.save()