#!/usr/bin/env python

import uuid
import datetime
import urllib2
import urllib
import httplib
import ssl
import socket
import json
import argparse
import sys

CURRENT_VERSION = "1.1-imanolbarba"
DEBUG = 0

EXTERNAL_IP_URL = "https://ipv4.wtfismyip.com/text"

TYPE_INFO = 0
TYPE_WARNING = 1
TYPE_ERROR = 2
TYPE_DEBUG = 3

BASH_YLW_TEXT = '\033[33m'
BASH_RED_TEXT = '\033[31m'
BASH_RESET_TEXT = '\033[0m'

#MY_API_KEY = "AAAAAAAAAAAAAAAA"

#Generates random UUID
def genUUID():
	return str(uuid.uuid4())

#Message logging function
#TODO: log to file
def msgLog(string,type = TYPE_INFO):
	time = str(datetime.datetime.now())

	if(type == TYPE_ERROR):
		print "[" + time + "]: " + BASH_RED_TEXT + "ERROR: " + string + BASH_RESET_TEXT

	elif(type == TYPE_WARNING):
		print "[" + time + "]: " + BASH_YLW_TEXT + "WARNING: " + string + BASH_RESET_TEXT
	
	elif(type == TYPE_INFO):
		print "[" + time + "]: " + "INFO: " + string		
	
	elif(type == TYPE_DEBUG):
		if(DEBUG):
			print "[" + time + "]: " + "DEBUG: " + string
	
	else: 
		print string

#External IP fetcher, expects plain text response
def fetchIP(): 
	text = urllib2.urlopen(EXTERNAL_IP_URL).read()
	return ''.join(text.split()) #remove whitespace

class APIConn:
	APIKey = ""
	username = ""
	
	def __init__(self):
		self.APIKey = MY_API_KEY

	def sendCmd(self,cmd,args):
		conn = httplib.HTTPSConnection(host="api.dreamhost.com", port=443)
		sock = socket.create_connection((conn.host, conn.port), conn.timeout, conn.source_address)
		conn.sock = ssl.wrap_socket(sock, conn.key_file, conn.cert_file, ssl_version=ssl.PROTOCOL_TLSv1)
		msgLog("Connection with Dreamhost established",TYPE_DEBUG)
		apiParams = { "key" : self.APIKey, "unique_id" : genUUID(), "cmd" : cmd, "format" : "json" }
		if(len(args) != 0):
			apiParams.update(args)
		params = urllib.urlencode(apiParams)
		msgLog("Sending request with parameters " + params,TYPE_DEBUG)
		conn.request("POST","/",params)
		answer = conn.getresponse().read()
		msgLog("Received answer: " + answer, TYPE_DEBUG)
		conn.close()
		msgLog("Connection closed", TYPE_DEBUG)
		return answer


def listDNS(apiconn):
	cmd = "dns-list_records"
	args = {}
	data = json.loads(apiconn.sendCmd(cmd,args))
	records = {}
	for i in xrange(len(data["data"])):
		if(data["data"][i]["type"] == "A"):
			records[str(data["data"][i]["record"])] = str(data["data"][i]["value"])
	msgLog("Returned records: " + str(records), TYPE_DEBUG)
	return records

def updateDNS(apiconn,record,oldIP,newIP):
	removeDNS(apiconn,record,oldIP)
	addDNS(apiconn,record,newIP)
	msgLog("Successfully updated " + record,TYPE_INFO)

def addDNS(apiconn,record,IP):
	cmd = "dns-add_record"
        args = {"record" : record , "type" : "A" , "value" : IP}
        data = json.loads(apiconn.sendCmd(cmd,args))
        if(data["result"] != "success"):
                if(data["data"] == "no_such_zone"):
                        msgLog("Server returned no_such_zone, this is a known issue in the Dreamhost API, go to your Panel and add a A record for your domain with the IP 0.0.0.0, then run " + sys.argv[0] + " --add " + record,TYPE_WARNING)

                msgLog("Error while adding record, server returned: " + data["data"],TYPE_ERROR)
                return
        msgLog("Successfully added " + record,TYPE_INFO)

def removeDNS(apiconn,record,IP):
	cmd = "dns-remove_record"
	args = {"record" : record , "type" : "A" , "value" : IP}
	data = json.loads(apiconn.sendCmd(cmd,args))
	if(data["result"] != "success"):
		msgLog("Error while removing old record, server returned: " + data["data"],TYPE_ERROR)
		return

	msgLog("Successfully removed " + record,TYPE_INFO)

if(MY_API_KEY == ""):
	msgLog("Missing API Key. Please open the script and add your key in the MY_API_KEY variable",TYPE_ERROR)
	exit()

api = APIConn()

parser = argparse.ArgumentParser(description="Update Dreamhost DNS records")
parser.add_argument('--verbose', action='store_true', dest="verbose", help="Enable debugging messages")
parser.add_argument('--version', action='version', version=str("DreamDDNS v" + CURRENT_VERSION))
group = parser.add_mutually_exclusive_group()
group.add_argument("--update", action="store", dest="update", help="Update DNS record RECORD", nargs=1)
group.add_argument("--list", action="store_true", dest="list", help="List current DNS records")
group.add_argument("--add", action="store", dest="add", help="Add DNS record RECORD", nargs=1)
group.add_argument("--remove", action="store", dest="remove", help="Remove DNS record RECORD", nargs=1)

if(len(sys.argv) == 1):
	parser.print_help()
	
argv = parser.parse_args()
if(argv.verbose):
	DEBUG = 1

if(argv.list):
	records = listDNS(api)
	for i in records:
		print i + ": " + records[i]
elif(argv.update != None):
	targetRecord = argv.update[0]
	newIP = fetchIP()
	records = listDNS(api)
	recordFound = False
	for i in records:
		if(i == targetRecord):
			recordFound = True
			if(records[i] == newIP):
				msgLog(targetRecord + " is already up-to-date. No further action needed",TYPE_INFO)
			else:
				msgLog("Updating " + targetRecord + "...",TYPE_INFO)
				updateDNS(api,targetRecord,records[i],newIP)

	if(recordFound == False):
		msgLog("Record not found",TYPE_ERROR)

elif(argv.add != None):
	targetRecord = argv.add[0]
	newIP = fetchIP()
	msgLog("Adding " + targetRecord + "...",TYPE_INFO)	
	addDNS(api,targetRecord,newIP)

elif(argv.remove != None):
	targetRecord = argv.remove[0]
	records = listDNS(api)
	recordFound = False
	for i in records:
		if(i == targetRecord):
			recordFound = True
			msgLog("Removing " + targetRecord + "...",TYPE_INFO)
			removeDNS(api,targetRecord,records[i])

	if(recordFound == False):
		msgLog("Record not found",TYPE_ERROR)