Add 'rebuild_db.py'

1 files changed, 691 insertions, 0 deletions

diff --git a/rebuild_db.py b/rebuild_db.py
new file mode 100644
index 0000000..ee42792
--- /dev/null
+++ b/rebuild_db.py
@@ -0,0 +1,691 @@
+#!/usr/bin/env python
+# -*- coding: iso-8859-1
+
+# 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 2 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, write to the Free Software
+# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
+
+__title__="KeyJ's iPod shuffle Database Builder patched for modern Python 3"
+__version__="1.0-rc2"
+__author__="Martin Fiedler"
+__email__="martin.fiedler@gmx.net"
+
+""" VERSION HISTORY
+1.0-rc1 (2006-04-26)
+    * finally(!) added the often-requested auto-rename feature (-r option)
+0.7-pre1 (2005-06-09)
+    * 0.6-final was skipped because of the huge load of new (experimental)
+      features in this version :)
+    * rule files allow for nice and flexible fine tuning
+    * fixed --nochdir and --nosmart bugs
+    * numerical sorting (i.e. "track2.mp3" < "track10.mp3")
+    * iTunesSD entries are now copied over from the old file if possible; this
+      should preserve the keys for .aa files
+0.6-pre2 (2005-05-02)
+    * fixed file type bug (thanks to Nowhereman)
+    * improved audio book support (thanks to Nowhereman)
+    * files called $foo.book.$type (like example.book.mp3) are stored as
+      audio books now
+    * the subdirectories of /iPod_Control/Music are merged as if they were a
+      single directory
+0.6-pre1 (2005-04-23)
+    * always starts from the directory of the executable now
+    * output logging
+    * generating iTunesPState and iTunesStats so that the iPod won't overwrite
+      iTunesShuffle anymore
+    * -> return of smart shuffle ;)
+    * directory display order is identical to playback ordernow
+    * command line options and help
+    * interactive mode, configurable playback volume, directory limitation
+0.5 (2005-04-15)
+    * major code refactoring (thanks to Andre Kloss)
+    * removed "smart shuffle" again -- the iPod deleted the file anyway :(
+    * common errors are now reported more concisely
+    * dot files (e.g. ".hidden_file") are now ignored while browsing the iPod
+0.4 (2005-03-20)
+    * fixed iPod crashes after playing the shuffle playlist to the end
+    * fixed incorrect databse entries for non-MP3 files
+0.3 (2005-03-18)
+    * Python version now includes a "smart shuffle" feature
+0.2 (2005-03-15)
+    * added Python version
+0.1 (2005-03-13)
+    * initial public release, Win32 only
+"""
+
+
+import sys,os,os.path,array,getopt,random,types,fnmatch,operator,string
+from functools import reduce
+
+KnownProps=('filename','size','ignore','type','shuffle','reuse','bookmark')
+Rules=[
+  ([('filename','~','*.mp3')],          {'type':1, 'shuffle':1, 'bookmark':0}),
+  ([('filename','~','*.m4?')],          {'type':2, 'shuffle':1, 'bookmark':0}),
+  ([('filename','~','*.m4b')],          {          'shuffle':0, 'bookmark':1}),
+  ([('filename','~','*.aa')],           {'type':1, 'shuffle':0, 'bookmark':1, 'reuse':1}),
+  ([('filename','~','*.wav')],          {'type':4, 'shuffle':0, 'bookmark':0}),
+  ([('filename','~','*.book.???')],     {          'shuffle':0, 'bookmark':1}),
+  ([('filename','~','*.announce.???')], {          'shuffle':0, 'bookmark':0}),
+  ([('filename','~','/recycled/*')],    {'ignore':1}),
+]
+
+Options={
+  "volume":None,
+  "interactive":False,
+  "smart":True,
+  "home":True,
+  "logging":True,
+  "reuse":1,
+  "logfile":"rebuild_db.log.txt",
+  "rename":False
+}
+domains=[]
+total_count=0
+KnownEntries={}
+
+
+################################################################################
+
+
+def open_log():
+  global logfile
+  if Options['logging']:
+    try:
+      logfile=open(Options['logfile'],"w")
+    except IOError:
+      logfile=None
+  else:
+   logfile=None
+
+
+def log(line="",newline=True):
+  global logfile
+  if newline:
+    print(line)
+    line+="\n"
+  else:
+    print(line, end=' ')
+    line+=" "
+  if logfile:
+    try:
+      logfile.write(line)
+    except IOError:
+      pass
+
+
+def close_log():
+  global logfile
+  if logfile:
+    logfile.close()
+
+
+def go_home():
+  if Options['home']:
+    try:
+      os.chdir(os.path.split(sys.argv[0])[0])
+    except OSError:
+      pass
+
+
+def filesize(filename):
+  try:
+    return os.stat(filename)[6]
+  except OSError:
+    return None
+
+
+################################################################################
+
+
+def MatchRule(props,rule):
+  try:
+    prop,op,ref=props[rule[0]],rule[1],rule[2]
+  except KeyError:
+    return False
+  if rule[1]=='~':
+    return fnmatch.fnmatchcase(prop.lower(),ref.lower())
+  elif rule[1]=='=':
+    return cmp(prop,ref)==0
+  elif rule[1]=='>':
+    return cmp(prop,ref)>0
+  elif rule[1]=='<':
+    return cmp(prop,ref)<0
+  else:
+    return False
+
+
+def ParseValue(val):
+  if len(val)>=2 and ((val[0]=="'" and val[-1]=="'") or (val[0]=='"' and val[-1]=='"')):
+    return val[1:-1]
+  try:
+    return int(val)
+  except ValueError:
+    return val
+
+def ParseRule(rule):
+  sep_pos=min([rule.find(sep) for sep in "~=<>" if rule.find(sep)>0])
+  prop=rule[:sep_pos].strip()
+  if not prop in KnownProps:
+    log("WARNING: unknown property `%s'"%prop)
+  return (prop,rule[sep_pos],ParseValue(rule[sep_pos+1:].strip()))
+
+def ParseAction(action):
+  prop,value=list(map(string.strip,action.split('=',1)))
+  if not prop in KnownProps:
+    log("WARNING: unknown property `%s'"%prop)
+  return (prop,ParseValue(value))
+
+def ParseRuleLine(line):
+  line=line.strip()
+  if not(line) or line[0]=="#":
+    return None
+  try:
+    # split line into "ruleset: action"
+    tmp=line.split(":")
+    ruleset=list(map(string.strip,":".join(tmp[:-1]).split(",")))
+    actions=dict(list(map(ParseAction,tmp[-1].split(","))))
+    if len(ruleset)==1 and not(ruleset[0]):
+      return ([],actions)
+    else:
+      return (list(map(ParseRule,ruleset)),actions)
+  except OSError: #(ValueError,IndexError,KeyError):
+    log("WARNING: rule `%s' is malformed, ignoring"%line)
+    return None
+  return None
+
+
+################################################################################
+
+
+def safe_char(c):
+  if c in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_":
+    return c
+  return "_"
+
+def rename_safely(path,name):
+  base,ext=os.path.splitext(name)
+  newname=''.join(map(safe_char,base))
+  if name==newname+ext:
+    return name
+  if os.path.exists("%s/%s%s"%(path,newname,ext)):
+    i=0
+    while os.path.exists("%s/%s_%d%s"%(path,newname,i,ext)):
+      i+=1
+    newname+="_%d"%i
+  newname+=ext
+  try:
+    os.rename("%s/%s"%(path,name),"%s/%s"%(path,newname))
+  except OSError:
+    pass  # don't fail if the rename didn't work
+  return newname
+
+
+def write_to_db(filename):
+  global iTunesSD,domains,total_count,KnownEntries,Rules
+
+  # set default properties
+  props={
+    'filename': filename,
+    'size': filesize(filename[1:]),
+    'ignore': 0,
+    'type': 1,
+    'shuffle': 1,
+    'reuse': Options['reuse'],
+    'bookmark': 0
+  }
+
+  # check and apply rules
+  for ruleset,action in Rules:
+    if reduce(operator.__and__,[MatchRule(props,rule) for rule in ruleset],True):
+      props.update(action)
+  if props['ignore']: return 0
+
+  # retrieve entry from known entries or rebuild it
+  entry=props['reuse'] and (filename in KnownEntries) and KnownEntries[filename]
+  if not entry:
+    header[29]=props['type']
+    # debug galt:
+    #galt = entry=header.tostring()
+    #print("typeofGALT",type(galt))
+    # this is apparently returning type bytes ?
+    galt = "".join([c+"\0" for c in filename[:261]])+ \
+      "\0"*(525-2*len(filename))
+    entry=header.tostring()+galt.encode() 
+
+# ORIG
+#    entry=header.tostring()+ \
+#      "".join([c+"\0" for c in filename[:261]])+ \
+#      "\0"*(525-2*len(filename))
+
+  # write entry, modifying shuffleflag and bookmarkflag at least
+
+  #galt = entry[557]
+  #print("typeofGALT",type(galt)," =",galt)
+
+  iTunesSD.write(entry[:555]+(chr(props['shuffle'])+chr(props['bookmark'])+chr(entry[557])).encode())
+#ORIG  iTunesSD.write(entry[:555]+chr(props['shuffle'])+chr(props['bookmark'])+entry[557])
+  if props['shuffle']: domains[-1].append(total_count)
+  total_count+=1
+  return 1
+
+
+def make_key(s):
+  if not s: return s
+  s=s.lower()
+  for i in range(len(s)):
+    if s[i].isdigit(): break
+  if not s[i].isdigit(): return s
+  for j in range(i,len(s)):
+    if not s[j].isdigit(): break
+  if s[j].isdigit(): j+=1
+  return (s[:i],int(s[i:j]),make_key(s[j:]))
+
+def key_repr(x):
+  if type(x)==tuple:
+    return "%s%d%s"%(x[0],x[1],key_repr(x[2]))
+  else:
+    return x
+
+#galt's attempt to substitute the cmp function removed
+def cmp(xx,yy):
+  x = key_repr(xx)
+  y = key_repr(yy)
+  if (x > y):
+    return 1
+  if (x < y):
+    return -1
+  return 0
+
+def cmp_key(a,b):
+  if type(a)==tuple and type(b)==tuple:
+    return cmp(a[0],b[0]) or cmp(a[1],b[1]) or cmp_key(a[2],b[2])
+  else:
+    return cmp(key_repr(a),key_repr(b))
+
+
+def file_entry(path,name,prefix=""):
+  if not(name) or name[0]==".": return None
+  fullname="%s/%s"%(path,name)
+  may_rename=not(fullname.startswith("./iPod_Control")) and Options['rename']
+  try:
+    if os.path.islink(fullname):
+      return None
+    if os.path.isdir(fullname):
+      if may_rename: name=rename_safely(path,name)
+      return (0,make_key(name),prefix+name)
+    if os.path.splitext(name)[1].lower() in (".mp3",".m4a",".m4b",".m4p",".aa",".wav"):
+      if may_rename: name=rename_safely(path,name)
+      return (1,make_key(name),prefix+name)
+  except OSError:
+    pass
+  return None
+
+
+# added by galt since it was not yet in python3.1.2
+def cmp_to_key(mycmp):
+    'Convert a cmp= function into a key= function'
+    class K(object):
+        def __init__(self, obj, *args):
+            self.obj = obj
+        def __lt__(self, other):
+            return mycmp(self.obj, other.obj) < 0
+        def __gt__(self, other):
+            return mycmp(self.obj, other.obj) > 0
+        def __eq__(self, other):
+            return mycmp(self.obj, other.obj) == 0
+        def __le__(self, other):
+            return mycmp(self.obj, other.obj) <= 0
+        def __ge__(self, other):
+            return mycmp(self.obj, other.obj) >= 0
+        def __ne__(self, other):
+            return mycmp(self.obj, other.obj) != 0
+    return K
+
+
+def browse(path, interactive):
+  global domains
+
+  if path[-1]=="/": path=path[:-1]
+  displaypath=path[1:]
+  if not displaypath: displaypath="/"
+
+  if interactive:
+    while 1:
+      try:
+        choice=input("include `%s'? [(Y)es, (N)o, (A)ll] "%displaypath)[:1].lower()
+      except EOFError:
+        raise KeyboardInterrupt
+      if not choice: continue
+      if choice in "at":    # all/alle/tous/<dontknow>
+        interactive=0
+        break
+      if choice in "yjos":  # yes/ja/oui/si
+        break
+      if choice in "n":     # no/nein/non/non?
+        return 0
+
+  try:
+    files=[_f for _f in [file_entry(path,name) for name in os.listdir(path)] if _f]
+  except OSError:
+    return
+
+  if path=="./iPod_Control/Music":
+    subdirs=[x[2] for x in files if not x[0]]
+    files=[x for x in files if x[0]]
+    for dir in subdirs:
+      subpath="%s/%s"%(path,dir)
+      try:
+        files.extend([x for x in [file_entry(subpath,name,dir+"/") for name in os.listdir(subpath)] if x and x[0]])
+      except OSError:
+        pass
+
+#orig  files.sort(cmp_key)
+  files.sort(key=cmp_to_key(cmp_key))
+  count=len([None for x in files if x[0]])
+  if count: domains.append([])
+
+  real_count=0
+  for item in files:
+    fullname="%s/%s"%(path,item[2])
+    if item[0]:
+      real_count+=write_to_db(fullname[1:])
+    else:
+      browse(fullname,interactive)
+
+  if real_count==count:
+    log("%s: %d files"%(displaypath,count))
+  else:
+    log("%s: %d files (out of %d)"%(displaypath,real_count,count))
+
+
+################################################################################
+
+
+def stringval(i):
+  if i<0: i+=0x1000000
+  return "%c%c%c"%(i&0xFF,(i>>8)&0xFF,(i>>16)&0xFF)
+
+def listval(i):
+  if i<0: i+=0x1000000
+  return [i&0xFF,(i>>8)&0xFF,(i>>16)&0xFF]
+
+
+def make_playback_state(volume=None):
+  # I'm not at all proud of this function. Why can't stupid Python make strings
+  # mutable?!
+  log("Setting playback state ...",False)
+  PState=[]
+  try:
+    f=open("iPod_Control/iTunes/iTunesPState","rb")
+    a=array.array('B')
+    a.frombytes(f.read())
+    PState=a.tolist()
+    f.close()
+  except IOError as EOFError:
+    del PState[:]
+  if len(PState)!=21:
+    PState=listval(29)+[0]*15+listval(1)  # volume 29, FW ver 1.0
+  PState[3:15]=[0]*6+[1]+[0]*5  # track 0, shuffle mode, start of track
+  if volume is not None:
+    PState[:3]=listval(volume)
+  try:
+    f=open("iPod_Control/iTunes/iTunesPState","wb")
+    array.array('B',PState).tofile(f)
+    f.close()
+  except IOError:
+    log("FAILED.")
+    return 0
+  log("OK.")
+  return 1
+
+
+def make_stats(count):
+  log("Creating statistics file ...",False)
+  try:
+    open("iPod_Control/iTunes/iTunesStats","wb").write(\
+         (stringval(count)+"\0"*3+(stringval(18)+"\xff"*3+"\0"*12)*count).encode() )
+  except IOError:
+    log("FAILED.")
+    return 0
+  log("OK.")
+  return 1
+
+
+################################################################################
+
+
+def smart_shuffle():
+  try:
+    slice_count=max(list(map(len,domains)))
+  except ValueError:
+    return []
+  slices=[[] for x in range(slice_count)]
+  slice_fill=[0]*slice_count
+
+  for d in range(len(domains)):
+    used=[]
+    if not domains[d]: continue
+    for n in domains[d]:
+      # find slices where the nearest track of the same domain is far away
+      metric=[min([slice_count]+[min(abs(s-u),abs(s-u+slice_count),abs(s-u-slice_count)) for u in used]) for s in range(slice_count)]
+      thresh=(max(metric)+1)/2
+      farthest=[s for s in range(slice_count) if metric[s]>=thresh]
+
+      # find emptiest slices
+      thresh=(min(slice_fill)+max(slice_fill)+1)/2
+      emptiest=[s for s in range(slice_count) if slice_fill[s]<=thresh if (s in farthest)]
+
+      # choose one of the remaining candidates and add the track to the chosen slice
+      s=random.choice(emptiest or farthest)
+      slices[s].append((n,d))
+      slice_fill[s]+=1
+      used.append(s)
+
+  # shuffle slices and avoid adjacent tracks of the same domain at slice boundaries
+  seq=[]
+  last_domain=-1
+  for slice in slices:
+    random.shuffle(slice)
+    if len(slice)>2 and slice[0][1]==last_domain:
+      slice.append(slice.pop(0))
+    seq+=[x[0] for x in slice]
+    last_domain=slice[-1][1]
+  return seq
+
+
+def make_shuffle(count):
+  random.seed()
+  if Options['smart']:
+    log("Generating smart shuffle sequence ...",False)
+    seq=smart_shuffle()
+  else:
+    log("Generating shuffle sequence ...",False)
+    seq=list(range(count))
+    random.shuffle(seq)
+  try:
+    open("iPod_Control/iTunes/iTunesShuffle","wb").write(("".join(map(stringval,seq))).encode())
+  except IOError:
+    log("FAILED.")
+    return 0
+  log("OK.")
+  return 1
+
+
+################################################################################
+
+
+def main(dirs):
+  global header,iTunesSD,total_count,KnownEntries,Rules
+  log("Welcome to %s, version %s"%(__title__,__version__))
+  log()
+
+  try:
+    f=open("rebuild_db.rules","r")
+    Rules+=[_f for _f in list(map(ParseRuleLine,f.read().split("\n"))) if _f]
+    f.close()
+  except IOError:
+    pass
+
+  if not os.path.isdir("iPod_Control/iTunes"):
+    log("""ERROR: No iPod control directory found!
+Please make sure that:
+ (*) this program's working directory is the iPod's root directory
+ (*) the iPod was correctly initialized with iTunes""")
+    sys.exit(1)
+
+  header=array.array('B')
+  iTunesSD=None
+  try:
+    iTunesSD=open("iPod_Control/iTunes/iTunesSD","rb")
+    header.fromfile(iTunesSD,51)
+    if Options['reuse']:
+      iTunesSD.seek(18)
+      entry=iTunesSD.read(558)
+      while len(entry)==558:
+        filename=str(entry[33::2]).split("\0",1)[0]
+#orig        filename=entry[33::2].split("\0",1)[0]
+        KnownEntries[filename]=entry
+        entry=iTunesSD.read(558)
+  except (IOError,EOFError):
+    pass
+  if iTunesSD: iTunesSD.close()
+
+  if len(header)==51:
+    log("Using iTunesSD headers from existing database.")
+    if KnownEntries:
+      log("Collected %d entries from existing database."%len(KnownEntries))
+  else:
+    del header[18:]
+    if len(header)==18:
+      log("Using iTunesSD main header from existing database.")
+    else:
+      del header[:]
+      log("Rebuilding iTunesSD main header from scratch.")
+      header.fromlist([0,0,0,1,6,0,0,0,18]+[0]*9)
+    log("Rebuilding iTunesSD entry header from scratch.")
+    header.fromlist([0,2,46,90,165,1]+[0]*20+[100,0,0,1,0,2,0])
+
+  log()
+  try:
+    iTunesSD=open("iPod_Control/iTunes/iTunesSD","wb")
+    header[:18].tofile(iTunesSD)
+  except IOError:
+    log("""ERROR: Cannot write to the iPod database file (iTunesSD)!
+Please make sure that:
+ (*) you have sufficient permissions to write to the iPod volume
+ (*) you are actually using an iPod shuffle, and not some other iPod model :)""")
+    sys.exit(1)
+  del header[:18]
+
+  log("Searching for files on your iPod.")
+  try:
+    if dirs:
+      for dir in dirs:
+        browse("./"+dir,Options['interactive'])
+    else:
+      browse(".",Options['interactive'])
+    log("%d playable files were found on your iPod."%total_count)
+    log()
+    log("Fixing iTunesSD header.")
+    iTunesSD.seek(0)
+    iTunesSD.write(("\0%c%c"%(total_count>>8,total_count&0xFF)).encode())
+    iTunesSD.close()
+  except IOError:
+    log("ERROR: Some strange errors occured while writing iTunesSD.")
+    log("       You may have to re-initialize the iPod using iTunes.")
+    sys.exit(1)
+
+  if make_playback_state(Options['volume'])* \
+     make_stats(total_count)* \
+     make_shuffle(total_count):
+    log()
+    log("The iPod shuffle database was rebuilt successfully.")
+    log("Have fun listening to your music!")
+  else:
+    log()
+    log("WARNING: The main database file was rebuilt successfully, but there were errors")
+    log("         while resetting the other files. However, playback MAY work correctly.")
+
+
+################################################################################
+
+
+def help():
+  print("Usage: %s [OPTION]... [DIRECTORY]..."%sys.argv[0])
+  print("""Rebuild iPod shuffle database.
+
+Mandatory arguments to long options are mandatory for short options too.
+  -h, --help         display this help text
+  -i, --interactive  prompt before browsing each directory
+  -v, --volume=VOL   set playback volume to a value between 0 and 38
+  -s, --nosmart      do not use smart shuffle
+  -n, --nochdir      do not change directory to this scripts directory first
+  -l, --nolog        do not create a log file
+  -f, --force        always rebuild database entries, do not re-use old ones
+  -L, --logfile      set log file name
+
+Must be called from the iPod's root directory. By default, the whole iPod is
+searched for playable files, unless at least one DIRECTORY is specified.""")
+
+
+def opterr(msg):
+  print("parse error:",msg)
+  print("use `%s -h' to get help"%sys.argv[0])
+  sys.exit(1)
+
+def parse_options():
+  try:
+    opts,args=getopt.getopt(sys.argv[1:],"hiv:snlfL:r",\
+              ["help","interactive","volume=","nosmart","nochdir","nolog","force","logfile=","rename"])
+  except getopt.GetoptError as message:
+    opterr(message)
+  for opt,arg in opts:
+    if opt in ("-h","--help"):
+      help()
+      sys.exit(0)
+    elif opt in ("-i","--interactive"):
+      Options['interactive']=True
+    elif opt in ("-v","--volume"):
+      try:
+        Options['volume']=int(arg)
+      except ValueError:
+        opterr("invalid volume")
+    elif opt in ("-s","--nosmart"):
+      Options['smart']=False
+    elif opt in ("-n","--nochdir"):
+      Options['home']=False
+    elif opt in ("-l","--nolog"):
+      Options['logging']=False
+    elif opt in ("-f","--force"):
+      Options['reuse']=0
+    elif opt in ("-L","--logfile"):
+      Options['logfile']=arg
+    elif opt in ("-r","--rename"):
+      Options['rename']=True
+  return args
+
+
+################################################################################
+
+
+if __name__=="__main__":
+  args=parse_options()
+  go_home()
+  open_log()
+  try:
+    main(args)
+  except KeyboardInterrupt:
+    log()
+    log("You decided to cancel processing. This is OK, but please note that")
+    log("the iPod database is now corrupt and the iPod won't play!")
+  close_log()