#!/usr/bin/env python

import sys, os, getopt, httplib, mimetypes, urllib, urllib2, re

ALLOWED_EXTENSIONS = ('jpeg', 'jpg', 'png', 'gif', 'bmp', 'tif', 'tiff', 'swf')
MAX_FILE_SIZE = 1536 # size in kB
HOST = 'imageshack.us'
SELECTOR = '/'
USER_AGENT = "Mozilla/5.0 (X11; U; Linux i686; en-US; rv:1.7.12) \
Gecko/20050922 Firefox/1.0.7 (Debian package 1.0.7-1)"
REFERRER = 'http://imageshack.us/'
FORM_MAX_FILE_SIZE = '3145728'
FORM_SUBMIT = 'host it!'
REMOVE_THUMB_BAR_FIELD = ('rembar', '1')
HINT = 'SHELL UPLOADER:'
EMAIL = ''


POSTDATA = "content=%(content)s&description=%(description)s&type=%(type)s&expiry=%(expiry)s&name=%(name)s&save=0&s=Submit+Post"
URL = "http://en.pastebin.ca/index.php"
URL_PATTERN = re.compile('http://en.pastebin.ca/\d+')


def get_file_string(response):
  """
  Returns image identification string from the provided HTTPResponse
  object containg server response for the upload request.
  Returns None when the identification string can't be obtained.
  """
  it = iter(response.read().split('\n'))
  try:
    while True:
      line = it.next()
      if line.find(HINT) != -1:
        return line.split(':')[1].strip()
  except StopIteration:
    return None

def get_file_contents(file_name):
  """
  Returns contents of the file identified by provided file_name
  parameter.
  May raise IOError exception in case of a miserable failure.
  """
  f = open(file_name);
  data = f.read()
  f.close()
  return data

def send_file(file_name, additional_fields):
  """
  Attempts to upload provided file to imageshack.us hosting server.
  Returns identification string of the uploaded image.
  None is returned when upload failed or the identification string can't
  be obtained.
  May raise IOException when provided file can't be accessed properly.
  """

  fields = [
    ('MAX_FILE_SIZE', FORM_MAX_FILE_SIZE),
    ('refer', ''),
    ('brand', ''),
    ('email', EMAIL),
    ('submit', FORM_SUBMIT)
    ]

  fields.extend(additional_fields)

  file_data = ('fileupload', file_name, get_file_contents(file_name))

  content_type, body = encode_multipart_formdata(fields, file_data)

  h = httplib.HTTPConnection(HOST)
  h.putrequest('POST', SELECTOR)
  h.putheader('Content-Type', content_type)
  h.putheader('Content-Length', str(len(body)))
  h.putheader('Referer', REFERRER)
  h.putheader('User-Agent', USER_AGENT)
  h.endheaders()
  h.send(body)

  return get_file_string(h.getresponse())


def pastebin(content, description="", name="", expiry="", type="raw"):
	payload = POSTDATA % {
		"content" : urllib.quote_plus(content),
		"description" : urllib.quote_plus(description),
		"type" : 1,
		"expiry" : urllib.quote_plus(expiry),
		"name" : urllib.quote_plus(name)
		}
	req = urllib2.Request(URL, payload, {
		"Referer": "http://en.pastebin.ca/",
		"Content-Type": "application/x-www-form-urlencoded",
		"User-Agent": "Mozilla/5.0 (X11; U; Linux i686; en-US; rv:1.7.10) Gecko/20050813 Epiphany/1.7.6",
		"Accept-Language": "en",
		"Accept-Charset": "utf-8",
		})

	return urllib2.urlopen(req).read()

def got_page(res):
	matches = URL_PATTERN.findall(res)
	if len(matches) > 0:
		print matches[0]

def encode_multipart_formdata(fields, file_data):
  """
  Returns string containing MIME encoded form fields and file.
  """
  BOUNDARY = '----------HappyHappy__bouNdaRY_$'
  CRLF = '\r\n'
  L = []
  # add additional form fields
  for (key, value) in fields:
    L.append('--' + BOUNDARY)
    L.append('Content-Disposition: form-data; name="%s"' % key)
    L.append('')
    L.append(value)
  # add file
  L.append('--' + BOUNDARY)
  L.append('Content-Disposition: form-data; name="%s"; filename="%s"' \
    % (file_data[0], file_data[1]))
  L.append('Content-Type: %s' % get_content_type(file_data[1]))
  L.append('')
  L.append(file_data[2])
  L.append('--' + BOUNDARY + '--')
  L.append('')
  body = CRLF.join(L)
  content_type = 'multipart/form-data; boundary=%s' % BOUNDARY

  return content_type, body

def get_content_type(filename):
  """
  Returns MIME type for the provided file name.
  """
  return mimetypes.guess_type(filename)[0] or 'application/octet-stream'

def image_details(id_string):
  """
  Returns a hash with all usable parts of the image id string.
  """
  details = {'host': HOST}
  details['node'], details['dir'], details['file'] = id_string.split('/')
  details['file_prefix'], details['file_extension'] = details['file'].split('.')
  return details

def get_direct_url(id_string):
  """
  Returns direct URL to the image identification string.
  """
  return 'http://%s.%s/%s' % (id_string.split('/')[0], HOST, id_string)

def take_screenshot(what):
  """
    Takes a screenshot of terminal or root. Uses imagemagick
  """
  if what == "terminal":
      status = os.system("import -window $WINDOWID  tmp.png") 
  else:
      status = os.system("import -window root tmp.png")
 

def rm_tmp():
  """
  Removes tmp.png
  """
  os.system("rm tmp.png")
  
def parse_options():
  """
  Parses command line options. Images and output descriptors lists
  are returned.
  """
  try:
    options, images = getopt.getopt(sys.argv[1:], 'trhb:')
    if len(options) == 0:
      # make Usage: as default
      options.append(['-h', ''])
  except getopt.GetoptError, e:
    print >> sys.stderr, 'Error: ', e
    usage()
    sys.exit(1)

  return options, images

def send_screenshot():
  """
  Uploads the screenshot to ImageShack.us
  """
  additional_fields = []
  id_string = send_file("tmp.png", additional_fields)
  if id_string != None:
    write_output(output, id_string)
  else:
    print >> sys.stderr, 'Failed to upload file: %s' % image

def close_files(output):
  """
  Closes all output file descriptors.
  """
  for op, fp in output:
    fp.close()

def usage():
  """
  Prints little usage reference.
  """
  print """
Usage: %s [OPTIONS] <image files>
  -h print this message
  -r screenshot the desktop
  -t screenshot the terminal
  -b hosts text file in pastebion.ca
  """ % sys.argv[0]

def main():
  """Main function, one to rule them all."""
  
  try:
    options, images = parse_options()
    additional_fields = []
    output = [] # holds pairs of output type symbols and their file descriptors
    image = "tmp.png"
    content = ""
    description = ""
    expiry = ""
    type = "raw" 
    name = ""


    for o, a in options:
      if o == '-t':
        take_screenshot("terminal")
      elif o == '-r':
        take_screenshot("root")
      elif o == '-h':
        usage()
        sys.exit()
      elif o == '-b':
        f = file(a, 'r')
        lines = None
        lines = f.readlines()
        content = ''.join(lines)
        f.close()
        res = pastebin(content, description, name, expiry, type)
        got_page(res)
        sys.exit()

    if image.split('.')[-1].lower() not in ALLOWED_EXTENSIONS:
        print >> sys.stderr, '%s: illegal file extension, skipping' % image

    if os.stat(image).st_size < MAX_FILE_SIZE * 1024:
        id_string = send_file(image, additional_fields)
        if id_string != None:
          print get_direct_url(id_string)
        else:
          print >> sys.stderr, 'Failed to upload file: %s' % image
    else:
        print >> sys.stderr, '%s: file to large, max. size is %s kilobytes.' \
          % (image, MAX_FILE_SIZE)
  except (IOError, OSError), e:
    print >> sys.stderr, '%s: %s' % (e.strerror, e.filename)
  else:
    # close all open files
    rm_tmp()

if __name__ == '__main__':
  main()