From b53e7bc62ce1ccc2294c0371d9676fabff8c8a32 Mon Sep 17 00:00:00 2001
From: altf4arnold <arnold.dech@adct.be>
Date: Sun, 2 Jun 2024 01:04:57 +0200
Subject: [PATCH] Final structure

---
 atlas/atlascreator.py | 16 +++++++++++++++-
 tldtester/admin.py    |  2 +-
 tldtester/models.py   |  2 ++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/atlas/atlascreator.py b/atlas/atlascreator.py
index 46a2854..b1820af 100644
--- a/atlas/atlascreator.py
+++ b/atlas/atlascreator.py
@@ -58,6 +58,7 @@ def webrequest(tld, stack):
         measurement = None
     dbwriter(tld, stack, measurement)
 
+
 def dbwriter(unicodetld, stack, measurement):
     tld = Atlas.objects.filter(unicodetld=unicodetld)
     tldstack = tld.filter(stack=stack)
@@ -70,6 +71,18 @@ def dbwriter(unicodetld, stack, measurement):
         db.stack = stack
     db.measurement = measurement
     db.save()
+    tld = TLD.objects.filter(unicodetld=unicodetld)
+    if tld.exists():
+        primary_key = tld.values_list('pk', flat=True).first()
+        db = TLD.objects.get(pk=primary_key)
+        if stack == 4:
+            db.atlasv4 = measurement
+            db.save()
+        elif stack == 6:
+            db.atlasv6 = measurement
+            db.save()
+        else:
+            print("Unknown IP version")
 
 
 def main():
@@ -83,5 +96,6 @@ def main():
         webrequest(tld, 4)
         webrequest(tld, 6)
 
+
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
diff --git a/tldtester/admin.py b/tldtester/admin.py
index e671b6c..fdba353 100644
--- a/tldtester/admin.py
+++ b/tldtester/admin.py
@@ -7,7 +7,7 @@ import threading
 
 
 class tlds(ExtraButtonsMixin, admin.ModelAdmin):
-    list_display = ('tld', 'nsamount', 'v4nsamount', 'v6nsamount', 'dnssec', 'lastEdition')
+    list_display = ('tld', 'nsamount', 'v4nsamount', 'v6nsamount', 'dnssec', 'atlasv4', 'atlasv6', 'lastEdition')
 
     @button(change_form=True, html_attrs={'style': 'background-color:#88FF88;color:black'})
     def refresh(self, request):
diff --git a/tldtester/models.py b/tldtester/models.py
index 7c759b1..81eeefd 100644
--- a/tldtester/models.py
+++ b/tldtester/models.py
@@ -39,6 +39,8 @@ class TLD(models.Model):
     organisation = models.CharField(max_length=100)
     rdap = models.CharField(max_length=10, default="No")
     link = models.CharField(max_length=800, default="https://tldtest.net/")
+    atlasv4 = models.IntegerField(default=0, blank=True, null=True)
+    atlasv6 = models.IntegerField(default=0, blank=True, null=True)
 
     def __str__(self):
         return self.tld