#!/usr/bin/python
"""'ndu' version 0.1 Alpha - retrieve the content length of a given URI
Copyright (C) 2007 beplacid.com/beplacid.net
Usage: ndu [OPTIONS] URL [URL URL...]
Options:
	-h : this message (help)
	-c : convert the content length to KBs/MBs/GBs
"""
    #This program is free software; you can redistribute it and/or modify
    #it under the terms of the GNU General Public License as published by
    #the Free Software Foundation; either version 3 of the License, or
    #(at your option) any later version.

    #This program is distributed in the hope that it will be useful,
    #but WITHOUT ANY WARRANTY; without even the implied warranty of
    #MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    #GNU General Public License for more details.

    #You should have received a copy of the GNU General Public License
    #along with this program.  If not, see <http://www.gnu.org/licenses/>.

	# Copyright (C) 2007 Alex Collins

import httplib
import urlparse
import socket
import getopt
import sys

global OPT_CONVERT
OPT_CONVERT = False# whether the output is in KBs/MBs/GBs

# standard data sizes in bytes
ONE_KB = 1024
ONE_MB = ONE_KB * ONE_KB
ONE_GB = ONE_MB * ONE_MB

def main():
	""" the wrapper for entire script """
	urls = parse_args(); 

	if len(urls) > 0:
		print get_url_size(urls)
	else:
		print __doc__
		sys.exit(1)

def get_url_size(urls):
	""" Retrieves the content-length from the given url """
	# now retrieve the url
	# parse it into a tuple: (scheme,domain loc,uri,params)
	try:
		ret_str = ''
		format = ' bytes'
		# iterate through our URLs
		for url in urls:
			ret_str += '%s:' % url
			parts = urlparse.urlparse(url)
			http_req = httplib.HTTPConnection(parts[1])
			http_req.request('GET',parts[2])# request the URI
			http_res = http_req.getresponse()
			if http_res.status == httplib.OK:
				uri_len = http_res.getheader('Content-length')
				if uri_len != None: uri_len = long(uri_len)
			else:
				# bad url perhaps, print a nice message
				print "%s: %s" % (urls,httplib.responses[http_res.status])
				sys.exit(5)

			# now convert the length
			if OPT_CONVERT:
				# convert to KBs/MBs/GBs
				if uri_len < ONE_MB:
					uri_len = float(uri_len / ONE_KB)
					format = 'KB'
				elif uri_len >= ONE_MB and uri_len < ONE_GB:
					uri_len = float(uri_len / ONE_MB)
					format = 'MB'
				elif uri_len >= ONE_GB:
					uri_len = float(uri_len / ONE_GB)
					format = 'GB'
			
			ret_str += " %d%s\n" % (uri_len,format)
	except socket.error, (msg):
		print "%s: %s" % (urls,msg[1])
		sys.exit(4)
	return ret_str

def parse_args():
	""" Parses the arguments given to the script. Returns the url if given, otherwise prints the relevant error and exits """
	try:
		opts, args = getopt.getopt(sys.argv[1:], "hc", ["help","convert"])
	except getopt.error:
		print "Error parsing arguments!"
		print __doc__
		sys.exit(1) 

	url = None
	if len(args) == 0:
		print "Please provide a URL"
		print __doc__
		sys.exit(2)

	for o,a in opts:
		if o in ('-h','--help'):
			# -h or --help
			print __doc__
			sys.exit(0)
		elif o in ('-c','--convert'):
			# convert the content length to readable format
			global OPT_CONVERT
			OPT_CONVERT = True
		else:
			print "Unknown option '%s'" % o
			print __doc__
			sys.exit(3)
	return args

if __name__ == '__main__':
	main()