From 84b121ace165595037997d55680bfabc9a3ed5fb Mon Sep 17 00:00:00 2001
From: altf4arnold <arnold.dech@adct.be>
Date: Sun, 2 Jun 2024 00:47:31 +0200
Subject: [PATCH] Should be functionnal

---
 atlas/atlascreator.py | 48 ++++++++++++++++++++++++++++++++++---------
 atlas/models.py       |  2 +-
 2 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/atlas/atlascreator.py b/atlas/atlascreator.py
index 40fd22b..dc4cb3e 100644
--- a/atlas/atlascreator.py
+++ b/atlas/atlascreator.py
@@ -1,8 +1,11 @@
 import requests
 import config
+from tldtester.models import TLD
+from .models import Atlas
 
 
-def webrequest():
+def webrequest(tld, stack):
+    description_string = ("DNS measurement for " + tld + " in IPv" + str(stack))
     url = "https://atlas.ripe.net/api/v2/measurements/"
     api_key = config.ATLAS_API
     headers = {
@@ -13,9 +16,9 @@ def webrequest():
         "definitions": [
             {
                 "type": "dns",
-                "af": 6,
+                "af": stack,
                 "resolve_on_probe": True,
-                "description": "DNS measurement for be",
+                "description": description_string,
                 "query_class": "IN",
                 "query_type": "SOA",
                 "use_macros": False,
@@ -32,7 +35,7 @@ def webrequest():
                 "timeout": 5000,
                 "use_probe_resolver": True,
                 "set_nsid_bit": True,
-                "query_argument": "be"
+                "query_argument": tld
             }
         ],
         "probes": [
@@ -48,12 +51,37 @@ def webrequest():
 
     response = requests.post(url, headers=headers, json=data)
 
-    print(response.status_code)
-    print(response.json())
-    print("https://atlas.ripe.net/measurementdetail/")
+    if response.status_code == 201:
+        data = (response.json())
+        measurement = data['measurements'][0]
+    else:
+        measurement = None
+    dbwriter(tld, stack, measurement)
+
+def dbwriter(unicodetld, stack, measurement):
+    tld = Atlas.objects.filter(unicodetld=unicodetld)
+    tldstack = tld.filter(stack=stack)
+    if tldstack.exists():
+        primary_key = tldstack.values_list('pk', flat=True).first()
+        db = Atlas.objects.get(pk=primary_key)
+    else:
+        db = Atlas()
+        db.unicodetld = unicodetld
+        db.stack = stack
+    db.measurement = measurement
+    db.save()
 
-def dbwriter(response):
-    print(response)
 
 def main():
-    dbwriter(webrequest())
\ No newline at end of file
+    unicodetlds = []
+    # This will get the TLD's in unicode format from the database and put them in the list
+    tlds = TLD.objects.all().order_by('tld')
+    for tld in tlds:
+        db = TLD.objects.get(tld=tld)
+        unicodetlds.append(db.unicodetld)
+    for tld in unicodetlds:
+        webrequest(tld, 4)
+        webrequest(tld, 6)
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/atlas/models.py b/atlas/models.py
index 332a0da..a1f16be 100644
--- a/atlas/models.py
+++ b/atlas/models.py
@@ -5,7 +5,7 @@ class Atlas(models.Model):
     STACK = ((0, "Unknown"), (4, "IPv4"), (6, "IPv6"))
     unicodetld = models.CharField(max_length=100)
     stack = models.IntegerField(default=0, choices=STACK)
-    measurement = models.IntegerField(default=0)
+    measurement = models.IntegerField(default=0, blank=True, null=True)
     lastEdition = models.DateTimeField(auto_now=True)
 
     def __str__(self):