package main

import (
	"fmt"
	"github.com/miekg/dns"
)

func lookupIP(hostname string, resolver string) []string {
	c := dns.Client{}
	m := dns.Msg{}
	m.SetQuestion(dns.Fqdn(hostname), dns.TypeA)
	r, _, err := c.Exchange(&m, resolver+":53")
	if err != nil {
		return nil
	}
	if r.Rcode != dns.RcodeSuccess {
		fmt.Println("Error: querying " + resolver + " for A record of " + hostname + " failed with code " + dns.RcodeToString[r.Rcode])
		return nil
	}
	ips := make([]string, 0)
	for _, ans := range r.Answer {
		if rec, ok := ans.(*dns.A); ok {
			ips = append(ips, rec.A.String())
		}
	}

	if len(ips) > 0 {
		return ips
	}

	c = dns.Client{}
	m = dns.Msg{}
	m.SetQuestion(dns.Fqdn(hostname), dns.TypeCNAME)
	r, _, err = c.Exchange(&m, resolver+":53")
	if err != nil {
		return nil
	}
	if r.Rcode != dns.RcodeSuccess {
		fmt.Println("Error: querying " + resolver + " for CNAME record of " + hostname + " failed with code " + dns.RcodeToString[r.Rcode])
		return nil
	}
	for _, ans := range r.Answer {
		if rec, ok := ans.(*dns.CNAME); ok {
			ips = lookupIP(rec.Target, resolver)
			for _, ip := range ips {
				ips = append(ips, ip)
			}
		}
	}

	if len(ips) > 0 {
		return ips
	}

	fmt.Println("Error: no IP addresses found for " + hostname)
	return nil
}