import socket
import json

# Sample DNS records (A, MX, and other types)
dns_records = {
    "gnomeradio.org": {
        "A": [
            "178.255.144.178",
        ],
        "AAAA": [],
        "CNAME": [],
        "MX": [
            "10 aamot.org.",
        ],
        "TXT": [],
        "SRV": []
    },
    "example.com": {
        "A": ["178.255.144.178"],
        "MX": [("piperpal.com", 10)],
        "TXT": ["piperpal.com"],
        "CNAME": ["www.piperpal.com"]
    },
    "test.com": {
        "A": ["178.255.144.178"],
        "MX": [("cdn.aasw.clh.no", 10)],
        "TXT": ["cdn.aasw.clh.no"],
        "CNAME": ["cdn.aasw.clh.no"]
    },
    "example.net": {
        "A": ["178.21.130.26"],
        "MX": [("mx.webhuset.no", 10)],
        "TXT": ["www.arkivarium.no"],
        "CNAME": ["www.arkivarium.no"]
    }
}

# Server configuration
HOST = '178.255.144.178'
PORT = 5453

# Function to handle DNS lookup requests
def handle_dns_lookup(data):
    domain = data.strip()
    if domain in dns_records:
        response_data = {
            "domain_name": domain,
            "dns_records": dns_records[domain]
        }
        return json.dumps(response_data)
    else:
        return json.dumps({"error": "Domain not found"})

# Function to handle DNS update requests
def handle_dns_update(data):
    try:
        request_data = json.loads(data)
        if request_data["action"] == "update":
            domain = request_data["domain"]
            record_type = request_data["type"]
            record_data = request_data["data"]
            if domain in dns_records:
                if record_type in dns_records[domain]:
                    dns_records[domain][record_type].append(record_data)
                else:
                    dns_records[domain][record_type] = [record_data]
            else:
                dns_records[domain] = {record_type: [record_data]}
            return json.dumps({"status": "success", "message": "DNS record updated successfully"})
        else:
            return json.dumps({"status": "error", "message": "Invalid action"})
    except Exception as e:
        return json.dumps({"status": "error", "message": f"Error: {e}"})

# Main server function
def run_server():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
        server_socket.bind((HOST, PORT))
        server_socket.listen()

        print(f"Server is listening on {HOST}:{PORT}")

        while True:
            conn, addr = server_socket.accept()
            with conn:
                print(f"Connected by {addr}")
                while True:
                    data = conn.recv(1024)
                    if not data:
                        break
                    data = data.decode()
                    if data.startswith("{") and data.endswith("}"):
                        response = handle_dns_update(data)
                    else:
                        response = handle_dns_lookup(data)
                    conn.sendall(response.encode())

if __name__ == "__main__":
    run_server()