import socket
import json

# Sample DNS records (A, MX, and other types)
dns_records = {
    "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 = '127.0.0.1'
PORT = 65432

# Function to handle DNS lookup requests
def handle_dns_lookup(data):
    domain = data.strip()
    if domain in dns_records:
        return json.dumps(dns_records[domain])
    else:
        return "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 "DNS record updated successfully"
        else:
            return "Invalid action"
    except Exception as e:
        return 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()