## Amazon CloudFront support
 ## Author: Michal Ludvig <michal@logix.cz>
 ##         http://www.logix.cz/michal
 ## License: GPL Version 2
 

 import sys

 import time
 import httplib
 from logging import debug, info, warning, error
 
 try:
 	import xml.etree.ElementTree as ET
 except ImportError:
 	import elementtree.ElementTree as ET
 

 from Config import Config
 from Exceptions import *

 from Utils import getTreeFromXml, appendXmlTextNode, getDictFromTree, dateS3toPython, sign_string, getBucketFromHostname, getHostnameFromBucket

 from S3Uri import S3Uri, S3UriS3

 
 def output(message):
 	sys.stdout.write(message + "\n")
 
 def pretty_output(label, message):
 	#label = ("%s " % label).ljust(20, ".")
 	label = ("%s:" % label).ljust(15)
 	output("%s %s" % (label, message))
 
 class DistributionSummary(object):
 	## Example:
 	##
 	## <DistributionSummary>
 	##	<Id>1234567890ABC</Id>
 	##	<Status>Deployed</Status>
 	##	<LastModifiedTime>2009-01-16T11:49:02.189Z</LastModifiedTime>
 	##	<DomainName>blahblahblah.cloudfront.net</DomainName>
 	##	<Origin>example.bucket.s3.amazonaws.com</Origin>
 	##	<Enabled>true</Enabled>
 	## </DistributionSummary>

 	def __init__(self, tree):
 		if tree.tag != "DistributionSummary":
 			raise ValueError("Expected <DistributionSummary /> xml, got: <%s />" % tree.tag)
 		self.parse(tree)
 
 	def parse(self, tree):
 		self.info = getDictFromTree(tree)
 		self.info['Enabled'] = (self.info['Enabled'].lower() == "true")
 

 	def uri(self):
 		return S3Uri("cf://%s" % self.info['Id'])
 

 class DistributionList(object):
 	## Example:
 	## 

 	## <DistributionList xmlns="http://cloudfront.amazonaws.com/doc/2010-06-01/">

 	##	<Marker />
 	##	<MaxItems>100</MaxItems>
 	##	<IsTruncated>false</IsTruncated>
 	##	<DistributionSummary>
 	##	... handled by DistributionSummary() class ...
 	##	</DistributionSummary>
 	## </DistributionList>
 
 	def __init__(self, xml):
 		tree = getTreeFromXml(xml)
 		if tree.tag != "DistributionList":
 			raise ValueError("Expected <DistributionList /> xml, got: <%s />" % tree.tag)
 		self.parse(tree)
 
 	def parse(self, tree):
 		self.info = getDictFromTree(tree)
 		## Normalise some items
 		self.info['IsTruncated'] = (self.info['IsTruncated'].lower() == "true")
 
 		self.dist_summs = []
 		for dist_summ in tree.findall(".//DistributionSummary"):
 			self.dist_summs.append(DistributionSummary(dist_summ))
 

 class Distribution(object):

 	## Example:
 	##

 	## <Distribution xmlns="http://cloudfront.amazonaws.com/doc/2010-06-01/">

 	##	<Id>1234567890ABC</Id>
 	##	<Status>InProgress</Status>
 	##	<LastModifiedTime>2009-01-16T13:07:11.319Z</LastModifiedTime>
 	##	<DomainName>blahblahblah.cloudfront.net</DomainName>
 	##	<DistributionConfig>
 	##	... handled by DistributionConfig() class ...
 	##	</DistributionConfig>
 	## </Distribution>
 
 	def __init__(self, xml):
 		tree = getTreeFromXml(xml)
 		if tree.tag != "Distribution":
 			raise ValueError("Expected <Distribution /> xml, got: <%s />" % tree.tag)
 		self.parse(tree)
 
 	def parse(self, tree):
 		self.info = getDictFromTree(tree)
 		## Normalise some items
 		self.info['LastModifiedTime'] = dateS3toPython(self.info['LastModifiedTime'])
 
 		self.info['DistributionConfig'] = DistributionConfig(tree = tree.find(".//DistributionConfig"))

 	
 	def uri(self):
 		return S3Uri("cf://%s" % self.info['Id'])

 
 class DistributionConfig(object):
 	## Example:
 	##
 	## <DistributionConfig>
 	##	<Origin>somebucket.s3.amazonaws.com</Origin>
 	##	<CallerReference>s3://somebucket/</CallerReference>
 	##	<Comment>http://somebucket.s3.amazonaws.com/</Comment>
 	##	<Enabled>true</Enabled>

 	##  <Logging>
 	##    <Bucket>bu.ck.et</Bucket>
 	##    <Prefix>/cf-somebucket/</Prefix>
 	##  </Logging>

 	## </DistributionConfig>
 

 	EMPTY_CONFIG = "<DistributionConfig><Origin/><CallerReference/><Enabled>true</Enabled></DistributionConfig>"

 	xmlns = "http://cloudfront.amazonaws.com/doc/2010-06-01/"

 	def __init__(self, xml = None, tree = None):
 		if not xml:
 			xml = DistributionConfig.EMPTY_CONFIG
 
 		if not tree:
 			tree = getTreeFromXml(xml)
 
 		if tree.tag != "DistributionConfig":
 			raise ValueError("Expected <DistributionConfig /> xml, got: <%s />" % tree.tag)
 		self.parse(tree)
 
 	def parse(self, tree):

 		self.info = getDictFromTree(tree)
 		self.info['Enabled'] = (self.info['Enabled'].lower() == "true")

 		if not self.info.has_key("CNAME"):
 			self.info['CNAME'] = []
 		if type(self.info['CNAME']) != list:
 			self.info['CNAME'] = [self.info['CNAME']]

 		self.info['CNAME'] = [cname.lower() for cname in self.info['CNAME']]

 		if not self.info.has_key("Comment"):
 			self.info['Comment'] = ""

 		## Figure out logging - complex node not parsed by getDictFromTree()
 		logging_nodes = tree.findall(".//Logging")
 		if logging_nodes:
 			logging_dict = getDictFromTree(logging_nodes[0])
 			logging_dict['Bucket'], success = getBucketFromHostname(logging_dict['Bucket'])
 			if not success:
 				warning("Logging to unparsable bucket name: %s" % logging_dict['Bucket'])
 			self.info['Logging'] = S3UriS3("s3://%(Bucket)s/%(Prefix)s" % logging_dict)
 		else:
 			self.info['Logging'] = None

 
 	def __str__(self):

 		tree = ET.Element("DistributionConfig")

 		tree.attrib['xmlns'] = DistributionConfig.xmlns
 
 		## Retain the order of the following calls!

 		appendXmlTextNode("Origin", self.info['Origin'], tree)
 		appendXmlTextNode("CallerReference", self.info['CallerReference'], tree)
 		for cname in self.info['CNAME']:

 			appendXmlTextNode("CNAME", cname.lower(), tree)

 		if self.info['Comment']:
 			appendXmlTextNode("Comment", self.info['Comment'], tree)

 		appendXmlTextNode("Enabled", str(self.info['Enabled']).lower(), tree)

 		if self.info['Logging']:
 			logging_el = ET.Element("Logging")
 			appendXmlTextNode("Bucket", getHostnameFromBucket(self.info['Logging'].bucket()), logging_el)
 			appendXmlTextNode("Prefix", self.info['Logging'].object(), logging_el)
 			tree.append(logging_el)

 		return ET.tostring(tree)

 
 class CloudFront(object):
 	operations = {

 		"CreateDist" : { 'method' : "POST", 'resource' : "" },
 		"DeleteDist" : { 'method' : "DELETE", 'resource' : "/%(dist_id)s" },

 		"GetList" : { 'method' : "GET", 'resource' : "" },
 		"GetDistInfo" : { 'method' : "GET", 'resource' : "/%(dist_id)s" },
 		"GetDistConfig" : { 'method' : "GET", 'resource' : "/%(dist_id)s/config" },
 		"SetDistConfig" : { 'method' : "PUT", 'resource' : "/%(dist_id)s/config" },
 	}
 
 	## Maximum attempts of re-issuing failed requests
 	_max_retries = 5
 
 	def __init__(self, config):
 		self.config = config
 
 	## --------------------------------------------------
 	## Methods implementing CloudFront API
 	## --------------------------------------------------
 
 	def GetList(self):
 		response = self.send_request("GetList")

 		response['dist_list'] = DistributionList(response['data'])
 		if response['dist_list'].info['IsTruncated']:
 			raise NotImplementedError("List is truncated. Ask s3cmd author to add support.")
 		## TODO: handle Truncated 
 		return response
 	

 	def CreateDistribution(self, uri, cnames_add = [], comment = None, logging = None):

 		dist_config = DistributionConfig()
 		dist_config.info['Enabled'] = True
 		dist_config.info['Origin'] = uri.host_name()
 		dist_config.info['CallerReference'] = str(uri)

 		if comment == None:
 			dist_config.info['Comment'] = uri.public_url()
 		else:
 			dist_config.info['Comment'] = comment
 		for cname in cnames_add:
 			if dist_config.info['CNAME'].count(cname) == 0:
 				dist_config.info['CNAME'].append(cname)

 		if logging:

 			dist_config.info['Logging'] = S3UriS3(logging)

 		request_body = str(dist_config)

 		debug("CreateDistribution(): request_body: %s" % request_body)
 		response = self.send_request("CreateDist", body = request_body)
 		response['distribution'] = Distribution(response['data'])

 		return response

 	def ModifyDistribution(self, cfuri, cnames_add = [], cnames_remove = [],

 	                       comment = None, enabled = None, logging = None):

 		if cfuri.type != "cf":
 			raise ValueError("Expected CFUri instead of: %s" % cfuri)
 		# Get current dist status (enabled/disabled) and Etag
 		info("Checking current status of %s" % cfuri)
 		response = self.GetDistConfig(cfuri)
 		dc = response['dist_config']
 		if enabled != None:
 			dc.info['Enabled'] = enabled
 		if comment != None:
 			dc.info['Comment'] = comment
 		for cname in cnames_add:
 			if dc.info['CNAME'].count(cname) == 0:
 				dc.info['CNAME'].append(cname)
 		for cname in cnames_remove:
 			while dc.info['CNAME'].count(cname) > 0:
 				dc.info['CNAME'].remove(cname)

 		if logging != None:

 			if logging == False:
 				dc.info['Logging'] = False
 			else:
 				dc.info['Logging'] = S3UriS3(logging)

 		response = self.SetDistConfig(cfuri, dc, response['headers']['etag'])
 		return response
 		

 	def DeleteDistribution(self, cfuri):
 		if cfuri.type != "cf":
 			raise ValueError("Expected CFUri instead of: %s" % cfuri)

 		# Get current dist status (enabled/disabled) and Etag
 		info("Checking current status of %s" % cfuri)
 		response = self.GetDistConfig(cfuri)
 		if response['dist_config'].info['Enabled']:
 			info("Distribution is ENABLED. Disabling first.")
 			response['dist_config'].info['Enabled'] = False
 			response = self.SetDistConfig(cfuri, response['dist_config'], 
 			                              response['headers']['etag'])
 			warning("Waiting for Distribution to become disabled.")
 			warning("This may take several minutes, please wait.")
 			while True:
 				response = self.GetDistInfo(cfuri)
 				d = response['distribution']
 				if d.info['Status'] == "Deployed" and d.info['Enabled'] == False:
 					info("Distribution is now disabled")
 					break
 				warning("Still waiting...")
 				time.sleep(10)
 		headers = {}
 		headers['if-match'] = response['headers']['etag']
 		response = self.send_request("DeleteDist", dist_id = cfuri.dist_id(),
 		                             headers = headers)
 		return response

 	
 	def GetDistInfo(self, cfuri):
 		if cfuri.type != "cf":
 			raise ValueError("Expected CFUri instead of: %s" % cfuri)
 		response = self.send_request("GetDistInfo", dist_id = cfuri.dist_id())
 		response['distribution'] = Distribution(response['data'])
 		return response

 	def GetDistConfig(self, cfuri):
 		if cfuri.type != "cf":
 			raise ValueError("Expected CFUri instead of: %s" % cfuri)
 		response = self.send_request("GetDistConfig", dist_id = cfuri.dist_id())
 		response['dist_config'] = DistributionConfig(response['data'])
 		return response
 	
 	def SetDistConfig(self, cfuri, dist_config, etag = None):
 		if etag == None:
 			debug("SetDistConfig(): Etag not set. Fetching it first.")
 			etag = self.GetDistConfig(cfuri)['headers']['etag']
 		debug("SetDistConfig(): Etag = %s" % etag)
 		request_body = str(dist_config)
 		debug("SetDistConfig(): request_body: %s" % request_body)
 		headers = {}
 		headers['if-match'] = etag
 		response = self.send_request("SetDistConfig", dist_id = cfuri.dist_id(),
 		                             body = request_body, headers = headers)
 		return response
 

 	## --------------------------------------------------
 	## Low-level methods for handling CloudFront requests
 	## --------------------------------------------------
 

 	def send_request(self, op_name, dist_id = None, body = None, headers = {}, retries = _max_retries):

 		operation = self.operations[op_name]

 		if body:
 			headers['content-type'] = 'text/plain'
 		request = self.create_request(operation, dist_id, headers)

 		conn = self.get_connection()

 		debug("send_request(): %s %s" % (request['method'], request['resource']))

 		conn.request(request['method'], request['resource'], body, request['headers'])
 		http_response = conn.getresponse()
 		response = {}
 		response["status"] = http_response.status
 		response["reason"] = http_response.reason
 		response["headers"] = dict(http_response.getheaders())
 		response["data"] =  http_response.read()
 		conn.close()
 
 		debug("CloudFront: response: %r" % response)
 

 		if response["status"] >= 500:

 			e = CloudFrontError(response)
 			if retries:
 				warning(u"Retrying failed request: %s" % op_name)
 				warning(unicode(e))
 				warning("Waiting %d sec..." % self._fail_wait(retries))
 				time.sleep(self._fail_wait(retries))
 				return self.send_request(op_name, dist_id, body, retries - 1)
 			else:
 				raise e
 
 		if response["status"] < 200 or response["status"] > 299:
 			raise CloudFrontError(response)
 
 		return response
 
 	def create_request(self, operation, dist_id = None, headers = None):
 		resource = self.config.cloudfront_resource + (
 		           operation['resource'] % { 'dist_id' : dist_id })
 
 		if not headers:
 			headers = {}
 
 		if headers.has_key("date"):
 			if not headers.has_key("x-amz-date"):
 				headers["x-amz-date"] = headers["date"]
 			del(headers["date"])
 		
 		if not headers.has_key("x-amz-date"):
 			headers["x-amz-date"] = time.strftime("%a, %d %b %Y %H:%M:%S +0000", time.gmtime())
 
 		signature = self.sign_request(headers)
 		headers["Authorization"] = "AWS "+self.config.access_key+":"+signature
 
 		request = {}
 		request['resource'] = resource
 		request['headers'] = headers
 		request['method'] = operation['method']
 
 		return request
 
 	def sign_request(self, headers):
 		string_to_sign = headers['x-amz-date']

 		signature = sign_string(string_to_sign)

 		debug(u"CloudFront.sign_request('%s') = %s" % (string_to_sign, signature))
 		return signature
 
 	def get_connection(self):
 		if self.config.proxy_host != "":
 			raise ParameterError("CloudFront commands don't work from behind a HTTP proxy")
 		return httplib.HTTPSConnection(self.config.cloudfront_host)
 
 	def _fail_wait(self, retries):
 		# Wait a few seconds. The more it fails the more we wait.
 		return (self._max_retries - retries + 1) * 3
 
 class Cmd(object):
 	"""
 	Class that implements CloudFront commands
 	"""

 	class Options(object):
 		cf_cnames_add = []
 		cf_cnames_remove = []
 		cf_comment = None
 		cf_enable = None

 		cf_logging = None

 
 		def option_list(self):
 			return [opt for opt in dir(self) if opt.startswith("cf_")]
 
 		def update_option(self, option, value):
 			setattr(Cmd.options, option, value)
 
 	options = Options()

 	dist_list = None
 
 	@staticmethod
 	def _get_dist_name_for_bucket(uri):
 		cf = CloudFront(Config())
 		debug("_get_dist_name_for_bucket(%r)" % uri)
 		assert(uri.type == "s3")
 		if Cmd.dist_list is None:
 			response = cf.GetList()
 			Cmd.dist_list = {}
 			for d in response['dist_list'].dist_summs:
 				Cmd.dist_list[getBucketFromHostname(d.info['Origin'])[0]] = d.uri()
 			debug("dist_list: %s" % Cmd.dist_list)
 		return Cmd.dist_list[uri.bucket()]
 
 	@staticmethod
 	def _parse_args(args):
 		cfuris = []
 		for arg in args:
 			uri = S3Uri(arg)
 			if uri.type == 's3':
 				try:
 					uri = Cmd._get_dist_name_for_bucket(uri)
 				except Exception, e:
 					debug(e)
 					raise ParameterError("Unable to translate S3 URI to CloudFront distribution name: %s" % uri)
 			if uri.type != 'cf':
 				raise ParameterError("CloudFront URI required instead of: %s" % arg)
 			cfuris.append(uri)
 		return cfuris

 	@staticmethod

 	def info(args):

 		cf = CloudFront(Config())

 		if not args:
 			response = cf.GetList()
 			for d in response['dist_list'].dist_summs:
 				pretty_output("Origin", S3UriS3.httpurl_to_s3uri(d.info['Origin']))
 				pretty_output("DistId", d.uri())
 				pretty_output("DomainName", d.info['DomainName'])
 				pretty_output("Status", d.info['Status'])
 				pretty_output("Enabled", d.info['Enabled'])
 				output("")
 		else:

 			cfuris = Cmd._parse_args(args)

 			for cfuri in cfuris:
 				response = cf.GetDistInfo(cfuri)
 				d = response['distribution']
 				dc = d.info['DistributionConfig']
 				pretty_output("Origin", S3UriS3.httpurl_to_s3uri(dc.info['Origin']))
 				pretty_output("DistId", d.uri())
 				pretty_output("DomainName", d.info['DomainName'])
 				pretty_output("Status", d.info['Status'])

 				pretty_output("CNAMEs", ", ".join(dc.info['CNAME']))
 				pretty_output("Comment", dc.info['Comment'])

 				pretty_output("Enabled", dc.info['Enabled'])

 				pretty_output("Logging", dc.info['Logging'] or "Disabled")

 				pretty_output("Etag", response['headers']['etag'])
 

 	@staticmethod
 	def create(args):
 		cf = CloudFront(Config())
 		buckets = []
 		for arg in args:
 			uri = S3Uri(arg)
 			if uri.type != "s3":
 				raise ParameterError("Bucket can only be created from a s3:// URI instead of: %s" % arg)
 			if uri.object():
 				raise ParameterError("Use s3:// URI with a bucket name only instead of: %s" % arg)
 			if not uri.is_dns_compatible():
 				raise ParameterError("CloudFront can only handle lowercase-named buckets.")
 			buckets.append(uri)
 		if not buckets:
 			raise ParameterError("No valid bucket names found")
 		for uri in buckets:
 			info("Creating distribution from: %s" % uri)

 			response = cf.CreateDistribution(uri, cnames_add = Cmd.options.cf_cnames_add, 

 			                                 comment = Cmd.options.cf_comment,
 			                                 logging = Cmd.options.cf_logging)

 			d = response['distribution']
 			dc = d.info['DistributionConfig']
 			output("Distribution created:")

 			pretty_output("Origin", S3UriS3.httpurl_to_s3uri(dc.info['Origin']))
 			pretty_output("DistId", d.uri())

 			pretty_output("DomainName", d.info['DomainName'])

 			pretty_output("CNAMEs", ", ".join(dc.info['CNAME']))
 			pretty_output("Comment", dc.info['Comment'])

 			pretty_output("Status", d.info['Status'])

 			pretty_output("Enabled", dc.info['Enabled'])

 			pretty_output("Etag", response['headers']['etag'])
 

 	@staticmethod
 	def delete(args):
 		cf = CloudFront(Config())

 		cfuris = Cmd._parse_args(args)

 		for cfuri in cfuris:
 			response = cf.DeleteDistribution(cfuri)
 			if response['status'] >= 400:
 				error("Distribution %s could not be deleted: %s" % (cfuri, response['reason']))

 			output("Distribution %s deleted" % cfuri)
 
 	@staticmethod
 	def modify(args):
 		cf = CloudFront(Config())

 		if len(args) > 1:

 			raise ParameterError("Too many parameters. Modify one Distribution at a time.")

 		try:
 			cfuri = Cmd._parse_args(args)[0]
 		except IndexError, e:
 			raise ParameterError("No valid Distribution URI found.")

 		response = cf.ModifyDistribution(cfuri,
 		                                 cnames_add = Cmd.options.cf_cnames_add,
 		                                 cnames_remove = Cmd.options.cf_cnames_remove,
 		                                 comment = Cmd.options.cf_comment,

 		                                 enabled = Cmd.options.cf_enable,
 		                                 logging = Cmd.options.cf_logging)

 		if response['status'] >= 400:
 			error("Distribution %s could not be modified: %s" % (cfuri, response['reason']))
 		output("Distribution modified: %s" % cfuri)
 		response = cf.GetDistInfo(cfuri)
 		d = response['distribution']
 		dc = d.info['DistributionConfig']
 		pretty_output("Origin", S3UriS3.httpurl_to_s3uri(dc.info['Origin']))
 		pretty_output("DistId", d.uri())
 		pretty_output("DomainName", d.info['DomainName'])
 		pretty_output("Status", d.info['Status'])
 		pretty_output("CNAMEs", ", ".join(dc.info['CNAME']))
 		pretty_output("Comment", dc.info['Comment'])
 		pretty_output("Enabled", dc.info['Enabled'])
 		pretty_output("Etag", response['headers']['etag'])