#!/usr/bin/env python

# ShackShag 0.1.2, 2007-07-04
# http://bachnik.com/projects/shackshag
# Author: Wiktor Bachnik, wiktor at bachnik dot com

# This software is released under GPL license:
# http://www.gnu.org/licenses/gpl.html
# Some inspiration and code snippets were taken from:
# http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/146306

# You're using this software at your own risk. The author shall not
# be held liable for ANY problems caused by using this software.
# Remember to read ImageShack Terms of Service before using their services
# It can be found at: http://reg.imageshack.us/content.php?page=rules

"""
ShackShag - imageshack batch image uploader.
"""

import sys, os, getopt, httplib, mimetypes

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 = ''

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)
  print body

  return get_file_string(h.getresponse())

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 get_thumbnail(id_string):
  """
  Returns thumbnail link for the image identification string.
  """
  return '<a href="http://%(node)s.%(host)s/my.php?image=%(file)s" \
target="_blank"><img src="\
http://%(node)s.%(host)s/%(node)s/%(dir)s/%(file_prefix)s.th.%(file_extension)s" \
border="0" alt="Free Image Hosting at www.ImageShack.us" /></a>' \
  % image_details(id_string)

def get_forum_link(id_string):
  """
  Returns forum hotlink code (full size) for the image identification string.
  """
  return '[url=http://%(host)s]\
[img=http://%(node)s.%(host)s/%(node)s/%(dir)s/%(file)s][/url]' \
  % image_details(id_string)

def write_output(files, id_string):
  """
  Writes image data to the specified output files.
  May rise IOError.
  """
  for op, fp in files:
    if op == '-s':
      print >> fp, get_direct_url(id_string)
    elif op == '-d':
      print >> fp, get_direct_url(id_string)
    elif op == '-t':
      print >> fp, get_thumbnail(id_string)
    elif op == '-f':
      print >> fp, get_forum_link(id_string)

def parse_options():
  """
  Parses command line options. Images and output descriptors lists
  are returned.
  """
  try:
    options, images = getopt.getopt(sys.argv[1:], 'hsrd:t:f:')
    if len(options) == 0:
      # make output to stdout default
      options.append(['-s', ''])
  except getopt.GetoptError, e:
    print >> sys.stderr, 'Error: ', e
    usage()
    sys.exit(1)

  return options, images

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
  -s use standard output (default when no output files are specified)
  -r remove size/resolution bar from thumbnails
  -d <direct links file>
  -t <thumbnails file>
  -f <forum links file>
  """ % 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
    for o, a in options:
      if o in ('-d', '-t', '-f'):
        output.append((o, open(a, 'w')))
      elif o == '-s':
        output.append((o, sys.stdout))
      elif o == '-r':
        additional_fields.append((REMOVE_THUMB_BAR_FIELD));
      elif o == '-h':
        usage()
        close_files(output)
        sys.exit()

    for image in images:
      if image.split('.')[-1].lower() not in ALLOWED_EXTENSIONS:
        print >> sys.stderr, '%s: illegal file extension, skipping' % image
        continue
      if os.stat(image).st_size < MAX_FILE_SIZE * 1024:
        id_string = send_file(image, additional_fields)
        if id_string != None:
          write_output(output, 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
    close_files(output)

if __name__ == '__main__':
  main()