From ca998f2be058536349cc4ccb3060d9a17ea42b6c Mon Sep 17 00:00:00 2001
From: crupest <crupest@outlook.com>
Date: Mon, 28 Nov 2022 17:35:11 +0800
Subject: Add aio dns.

---
 tool/modules/config.py |  5 +++--
 tool/modules/dns.py    | 42 ++++++++++++++++++++++++++++++++++++++++++
 tool/modules/nginx.py  |  4 ++++
 3 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 tool/modules/dns.py

(limited to 'tool/modules')

diff --git a/tool/modules/config.py b/tool/modules/config.py
index 28b09a3..b9ad818 100644
--- a/tool/modules/config.py
+++ b/tool/modules/config.py
@@ -1,18 +1,19 @@
 import pwd
 import grp
 import os
+import typing
 from rich.prompt import Prompt
 from .path import config_file_path
 
 
 class ConfigVar:
-    def __init__(self, name: str, description: str, default_value_generator, /, default_value_for_ask=None):
+    def __init__(self, name: str, description: str, default_value_generator: typing.Callable[[], str] | str, /, default_value_for_ask=str | None):
         """Create a config var.
 
         Args:
             name (str): The name of the config var.
             description (str): The description of the config var.
-            default_value_generator (typing.Callable([], str) | str): The default value generator of the config var. If it is a string, it will be used as the input prompt and let user input the value.
+            default_value_generator (typing.Callable[[], str] | str): The default value generator of the config var. If it is a string, it will be used as the input prompt and let user input the value.
         """
         self.name = name
         self.description = description
diff --git a/tool/modules/dns.py b/tool/modules/dns.py
new file mode 100644
index 0000000..0334353
--- /dev/null
+++ b/tool/modules/dns.py
@@ -0,0 +1,42 @@
+from os.path import *
+from io import StringIO
+import re
+from .nginx import *
+
+
+def generate_dns_zone(domain: str, ip: str, /, ttl: str | int = 600, *, enable_mail: bool = True, dkim: str | None = None) -> str:
+    result = f"$ORIGIN {domain}.\n\n"
+    result += "; A records\n"
+    result += f"@ {ttl} IN A {ip}\n"
+    subdomains = list_subdomain_names()
+    for subdomain in subdomains:
+        result += f"{subdomain} {ttl} IN A {ip}\n"
+
+    if enable_mail:
+        result += "\n; MX records\n"
+        result += f"@ {ttl} IN MX 10 mail.{domain}.\n"
+        result += "\n; SPF record\n"
+        result += f"@ {ttl} IN TXT \"v=spf1 mx ~all\"\n"
+        if dkim is not None:
+            result += "\n; DKIM record\n"
+            result += f"mail._domainkey {ttl} IN TEXT \"{dkim}\""
+            result += "\n; DMARC record\n"
+            result += "_dmarc {ttl} IN TXT \"v=DMARC1; p=none; rua=mailto:dmarc.report@{domain}; ruf=mailto:dmarc.report@{domain}; sp=none; ri=86400\"\n"
+    return result
+
+
+def get_dkim_from_mailserver(domain: str) -> str | None:
+    dkim_path = join(data_dir, "dms/config/opendkim/keys", domain, "mail.txt")
+    if not exists(dkim_path):
+        return None
+
+    buffer = StringIO()
+    subprocess.run(["sudo", "cat", dkim_path], stdout=buffer, check=True)
+    value = ""
+    for match in re.finditer("\"(.*)\"", str(buffer)):
+        value += match.groups[1]
+    return value
+
+
+def generate_dns_zone_with_dkim(domain: str, ip: str, /, ttl: str | int = 600) -> str:
+    return generate_dns_zone(domain, ip, ttl, enable_mail=True, dkim=get_dkim_from_mailserver(domain))
diff --git a/tool/modules/nginx.py b/tool/modules/nginx.py
index 087422b..4fc31cc 100755
--- a/tool/modules/nginx.py
+++ b/tool/modules/nginx.py
@@ -38,6 +38,10 @@ nginx_var_set = set.union(root_template.var_set,
                           static_file_template.var_set, reverse_proxy_template.var_set)
 
 
+def list_subdomain_names() -> list:
+    return [s["subdomain"] for s in server["sites"]]
+
+
 def list_subdomains(domain: str) -> list:
     return [f"{s['subdomain']}.{domain}" for s in server["sites"]]
 
-- 
cgit v1.2.3