#!/usr/bin/env python
import sys, os.path, glob, zipfile, platform, xml.etree.ElementTree

# publishers who have forsaken DRM
good = ['Tom Doherty']

if map(int, platform.mac_ver()[0].split('.')) > [10, 8]:
  bookdir = os.path.expanduser('~/Library/Containers/com.apple.BKAgentService/Data/Documents/iBooks')
else:
  bookdir = os.path.expanduser('~/Music/iTunes/iTunes Music/Books')

os.chdir(bookdir)

ok =  '\033[1;32mDRM-free    \033[0m'
bad = '\033[1;31mDRM-infested\033[0m'
#ok =  '<span style="color: green;">DRM-free    </span>'
#bad = '<span style="color: red;">DRM-infested</span>'

count = 0
salvageable = 0

def extract(meta):
  creator = ''
  status = ''
  pub = ''
  et = xml.etree.ElementTree.fromstring(meta)
  try:
    creator = et.findall('*{http://purl.org/dc/elements/1.1/}creator')
    creator = creator[0].text
    title = et.findall('*{http://purl.org/dc/elements/1.1/}title')
    title = title[0].text
  except:
    assert '!DOCTYPE plist' in meta
    next_tag = None
    for e in et[0].iter():
      if e.tag == 'key' and e.text in ('artistName', 'itemName'):
        next_tag = e.text
        continue
      if next_tag == 'artistName':
        creator = e.text
        next_tag = None
        continue
      elif next_tag == 'itemName':
        title = e.text
        next_tag = None
        continue
  pub = [x for x in good if x in meta]
  return creator, title, pub

def find_meta(file_list, opener):
  for m in file_list:
    if m.endswith('.opf') or m == 'iTunesMetadata.plist':
      meta = opener(m).read()
      return extract(meta)
  
for fn in glob.glob('*/*.epub'):
  status = ok
  suffix = ''
  if os.path.isdir(fn):
    suffix = '(directory)'
    if os.path.exists(fn + '/META-INF/encryption.xml'):
      status = bad
      count += 1
    meta = find_meta(os.listdir(fn), lambda x: open(fn + '/' + x))
  else:
    z = zipfile.ZipFile(fn)
    try:
      i = z.getinfo('META-INF/encryption.xml')
      status = bad
    except KeyError:
      pass
    meta = find_meta(z.namelist(), z.open)
    z.close()
  creator, title, pub = meta
  print status, fn, suffix
  print '\t', creator
  print '\t', title
  if status == bad and pub:
    print '\t\033[1;32mThis is published by', pub[0],
    #print '\t<span style="color: green;">This is published by', pub[0],
    print 'and could be re-downloaded DRM-free\033[0m'
    salvageable += 1

print count, 'books are DRM-infested'
print salvageable, 'could be cured'