From a0a099875f98124669776c38366ceb6e806732f2 Mon Sep 17 00:00:00 2001
From: Buddy Sandidge <buddy@sandidge.lol>
Date: Tue, 22 Nov 2016 22:49:01 -0800
Subject: [PATCH] Can pass in arguments from the command line arguments
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

 • Can set A and TXT records and set manual values
 • Can manually set content for records
 • Can pass in urls for getting IPv4 or IPv6 addresses
---
 main.go | 86 ++++++++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 73 insertions(+), 13 deletions(-)

diff --git a/main.go b/main.go
index 87aac34..c588e97 100644
--- a/main.go
+++ b/main.go
@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 	"os"
+	"strings"
 	"sync"
 
 	"git.xbudex.com/buddy/update-dns/dns"
@@ -12,6 +13,27 @@ import (
 	"github.com/urfave/cli"
 )
 
+var validTypes = []string{"A", "AAAA", "TXT"}
+
+type typeFlag struct {
+	value string
+}
+
+func (f *typeFlag) Set(v string) error {
+	value := strings.ToUpper(v)
+	for _, val := range validTypes {
+		if val == value {
+			f.value = val
+			return nil
+		}
+	}
+	return errors.Errorf("invalid type %s, expected %v", v, validTypes)
+}
+
+func (f *typeFlag) String() string {
+	return f.value
+}
+
 func main() {
 	app := cli.NewApp()
 	app.Name = "update-dns"
@@ -20,10 +42,28 @@ func main() {
 	app.Action = action
 
 	app.Flags = []cli.Flag{
+		cli.GenericFlag{
+			Name:  "type",
+			Value: &typeFlag{value: "AAAA"},
+			Usage: "type of dns record to create",
+		},
+
 		cli.StringFlag{
-			Name:  "url",
+			Name:  "content",
+			Value: "",
+			Usage: "content to set manually, defaults to fetching ip address if not set",
+		},
+
+		cli.StringFlag{
+			Name:  "get-ipv6-url",
 			Value: "https://whatismyipv6.buddy.wtf/json",
-			Usage: "url to use to get public ip address",
+			Usage: "url to use to get public IPv6 address",
+		},
+
+		cli.StringFlag{
+			Name:  "get-ipv4-url",
+			Value: "https://whatismyipv4.buddy.wtf/json",
+			Usage: "url to use to get public IPv4 address",
 		},
 
 		cli.StringFlag{
@@ -49,21 +89,42 @@ func action(context *cli.Context) error {
 	}
 
 	host := context.Args().Get(0)
+	kind := context.Generic("type").(*typeFlag).value
 
 	wg.Add(1)
 	ipClient := &ip.Client{
-		URLv4: "https://whatismyipv4.buddy.wtf/json",
-		URLv6: "https://whatismyipv6.buddy.wtf/json",
+		URLv4: context.String("get-ipv4-url"),
+		URLv6: context.String("get-ipv6-url"),
 	}
 
 	var getIPError error
-	var ip string
+	var content string
 	go func() {
 		defer wg.Done()
-		if result, err := ipClient.GetIPv6(); err != nil {
+
+		var err error
+		result := context.String("content")
+
+		switch kind {
+		case "A":
+			if result == "" {
+				result, err = ipClient.GetIPv4()
+			}
+		case "AAAA":
+			if result == "" {
+				result, err = ipClient.GetIPv6()
+			}
+		default:
+			if result == "" {
+				err = errors.Errorf("could not get content for type %s", kind)
+				return
+			}
+		}
+
+		if err != nil {
 			getIPError = errors.Wrap(err, "could not get IP address")
 		} else {
-			ip = result
+			content = result
 		}
 	}()
 
@@ -74,8 +135,7 @@ func action(context *cli.Context) error {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		fmt.Println("Getting records")
-		if results, err := client.GetRecords(host, "AAAA"); err != nil {
+		if results, err := client.GetRecords(host, kind); err != nil {
 			getRecordsError = errors.Wrap(err, "could not get records for host")
 		} else {
 			records = results
@@ -92,16 +152,16 @@ func action(context *cli.Context) error {
 		return errors.Wrap(getRecordsError, "could no get record")
 	}
 
-	if len(records) == 1 && records[0].Content == ip {
-		fmt.Printf("%s record for %s already set as %s\n", "AAAA", host, ip)
+	if len(records) == 1 && records[0].Content == content {
+		fmt.Printf("%s record for %s already set as %s\n", kind, host, content)
 		return nil
 	}
 
-	if err := client.CreateRecord(host, dns.Record{Type: "AAAA", Content: ip}); err != nil {
+	if err := client.CreateRecord(host, dns.Record{Type: kind, Content: content}); err != nil {
 		return errors.Wrap(err, "could not create record")
 	}
 
-	fmt.Printf("%s record created for %s as %s\n", "AAAA", host, ip)
+	fmt.Printf("%s record created for %s as %s\n", kind, host, content)
 
 	var deleteWG sync.WaitGroup
 	for _, record := range records {