Split out getSite into separate function

main.go

@@ -53,30 +53,29 @@ type HatsSite struct {
 	ScreenshotURL template.URL
 	Title         string
 	Notes         template.HTML
+	Error         template.HTML
 }
 
-func getSites(largest int, wd selenium.WebDriver) (sites []HatsSite, err error) {
+func getSite(domainName string, wd selenium.WebDriver) (hatsSite HatsSite) {
-	// TODO: 1hat.com 0hats.com, hats.com; possibly onehat.com, twohats.com, etc
-	for i := 2; i <= largest; i++ {
-		hatsSite := HatsSite{
-			DomainName: fmt.Sprintf("%dhats.com", i),
-			FetchTime:  time.Now(),
-		}
 	log.Printf("Retrieving info for %v\n", hatsSite.DomainName)
 
+	hatsSite.DomainName = domainName
+	hatsSite.FetchTime = time.Now()
+
 	// Check if domain is registered
 	query_result, err := whois.Whois(hatsSite.DomainName)
 	if err != nil {
-			return sites, err
+		hatsSite.Error = template.HTML(err.Error())
+		return
 	}
 	result, err := whoisparser.Parse(query_result)
 	if err == whoisparser.ErrNotFoundDomain {
 		hatsSite.Available = true
 		hatsSite.Notes = template.HTML(fmt.Sprintf("Register at <a href='https://www.namecheap.com/domains/registration/results/?domain=%v'>NameCheap</a>", hatsSite.DomainName))
-			sites = append(sites, hatsSite)
+		return
-			continue
 	} else if err != nil {
-			return sites, err
+		hatsSite.Error = template.HTML(err.Error())
+		return
 	}
 	hatsSite.Available = false
 	hatsSite.DomainInfo = result.Domain
@@ -86,30 +85,42 @@ func getSites(largest int, wd selenium.WebDriver) (sites []HatsSite, err error)
 	_, err = net.LookupHost(hatsSite.DomainName)
 	if err != nil {
 		hatsSite.Notes = template.HTML(fmt.Sprintf("DNS Error: <code>%v</code>", err.Error()))
-			sites = append(sites, hatsSite)
+		return
-			continue
 	}
 
 	// Get web page, take screenshot
 	err = wd.Get(fmt.Sprintf("http://%v/", hatsSite.DomainName))
 	if err != nil {
-			return sites, err
+		hatsSite.Error = template.HTML(err.Error())
+		return
 	}
 
 	hatsSite.Title, err = wd.Title()
 	if err != nil {
-			return sites, err
+		hatsSite.Error = template.HTML(err.Error())
+		return
 	}
 
 	screenshot, err := wd.Screenshot()
 	if err != nil {
-			return sites, err
+		hatsSite.Error = template.HTML(err.Error())
+		return
 	}
 	hatsSite.ScreenshotURL = template.URL(fmt.Sprintf("data:image/png;base64,%v", base64.StdEncoding.EncodeToString(screenshot)))
-
+	return
-		sites = append(sites, hatsSite)
 }
-	return sites, nil
+
+func getSites(largest int, wd selenium.WebDriver) (sites []HatsSite) {
+	domains := []string{"0hats.com", "1hat.com"} // Special cases
+	//domains = append(domains, "hats.com", "onehat.com", "twohats.com", "nhats.com") // Possibly more special cases?
+	for i := 2; i <= largest; i++ {
+		domains = append(domains, fmt.Sprintf("%dhats.com", i))
+	}
+
+	for _, domain := range domains {
+		sites = append(sites, getSite(domain, wd))
+	}
+	return sites
 }
 
 //go:embed template.html
@@ -171,11 +182,7 @@ func main() {
 		if err != nil {
 			panic(err)
 		}
-		sites, err := getSites(*largest, wd)
+		err = generateHTML(getSites(*largest, wd), file)
-		if err != nil {
-			fmt.Println(err)
-		}
-		err = generateHTML(sites, file)
 		if err != nil {
 			fmt.Println(err)
 		}