tvheadend/support/mkbundle

#!/usr/bin/env python
#
# Replacement for old mkbundle script that creates a full file heirarchy
#

import os, sys, re
import gzip
try:
  from cStringIO import StringIO as BytesIO
except:
  from io import BytesIO
from optparse import OptionParser

# Add reverse path split
def rsplit ( p ):
  i = p.find(os.path.sep)
  if i != -1:
    return (p[:i], p[i+1:])
  else:
    return (p, '')

# Process command line
optp = OptionParser()
optp.add_option('-z', '--gzip', action='store_true',
                help='Compress the files with gzip')
optp.add_option('-l', '--gzlevel', type='int', default=9,
                help='Specify compression level if using gzip')
optp.add_option('-o', '--output', default=None,
                help='Specify output file (default /dev/stdout)')
optp.add_option('-d', '--deps', default=None,
                help='Specify deps file to update during run')
(opts, args) = optp.parse_args()

# Setup
outf = sys.stdout
if opts.output:
  outf = open(opts.output, 'w')
depf = None
if opts.deps:
  depf = open(opts.deps, 'w')
  depf.write('%s: \\\n' % opts.output)

# Build heirarchy
root  = os.path.abspath(os.path.join(os.path.dirname(sys.argv[0]), '..'))
ents  = {}
for path in args:
  for (p, ds, fs) in os.walk(path, followlinks=True):
    p = os.path.abspath(p)
    n = p.replace(root+'/', '')
    t = ents
    while True:
      (d,n) = rsplit(n)
      if d.startswith('.'):
        fs = []
        break
      if d not in t:
        t[d] = {}
      t = t[d]
      if not n: break
    for f in fs:
      if f.startswith('.'): continue
      t[f] = None

# Output a file
def output_file ( path, name, idx, next = -1 ):
  n = 'NULL'
  if next >= 0: n = '&filebundle_entry_%06d' % next
  p = os.path.join(root, path, name);

  # Dep file
  if depf:
    depf.write('  %s\\\n' % p)

  # First the data
  outf.write('/* FILE : %s %s %d %d */\n' % (path, name, idx, next))
  outf.write('static const uint8_t filebundle_data_%06d[] = {' % idx)
  o = -1
  d = open(p, 'rb').read()
  if opts.gzip:
    o = len(d)
    l = opts.gzlevel
    t = BytesIO()
    if l < 0: l = 1
    if l > 9: l = 9
    z = gzip.GzipFile(filename=name, mode='w', compresslevel=l, fileobj=t)
    z.write(d)
    z.close()
    d = t.getvalue()
    t.close()
  i = 0
  for b in d:
    if not (i % 12): outf.write('\n  ')
    if type(b) == str:
      b = ord(b)
    outf.write('0x%02x,' % b)
    i = i + 1
  outf.write('\n')
  outf.write('};\n')
  
  outf.write('static filebundle_entry_t filebundle_entry_%06d = {\n' % idx)
  outf.write('  .type    = FB_FILE,\n')
  outf.write('  .name    = "%s",\n'  % name)
  outf.write('  .next    = %s,\n' % n)
  outf.write('  {\n')
  outf.write('    .f.size  = %d,\n' % len(d))
  outf.write('    .f.orig  = %d,\n' % o)
  outf.write('    .f.data  = filebundle_data_%06d\n' % idx)
  outf.write('  },\n')
  outf.write('};\n')
  outf.write('\n')

# Output a directory
def output_dir ( path, name, idx, child, count, next = -1 ):
  n = 'NULL'
  if next >= 0: n = '&filebundle_entry_%06d' % next
  outf.write('/* DIR: %s %s %d %d %d %d */\n' \
             % (path, name, idx, child, count, next))
  outf.write('static filebundle_entry_t filebundle_entry_%06d = {\n' % idx)
  outf.write('  .type    = FB_DIR,\n')
  outf.write('  .name    = "%s",\n'  % name)
  outf.write('  .next    = %s,\n' % n)
  outf.write('  {\n')
  outf.write('    .d.count = %d,\n' % count)
  outf.write('    .d.child = &filebundle_entry_%06d\n' % child)
  outf.write('  },\n')
  outf.write('};\n')
  outf.write('\n')

# Create output
def add_entry ( ents, path = "", name = "", idx = -1, next = -1 ):

  # Add children
  d = os.path.join(path, name)
  p = -1
  c = 0
  for k in ents:
    
    # File
    if ents[k] is None:
      output_file(d, k, idx+1, p)
      p = idx = idx + 1
      c = c + 1
      
    # Directory
    else:
      tmp = add_entry(ents[k], d, k, idx, p)
      if tmp != idx:
        p = idx = tmp
        c = c + 1

  # Add directory
  if p >= 0:
    if name:
      output_dir(path, name, idx+1, p, c, next)
      idx = idx + 1

  return idx

# Output header
outf.write('/* Auto-generated - DO NOT EDIT */\n')
outf.write('/* COMMAND: [%s] */\n' % (' '.join(sys.argv)))
outf.write('\n')
outf.write('#include "filebundle.h"\n')
outf.write('\n')

# Output entries
idx = add_entry(ents)

# Output top link
outf.write('filebundle_entry_t *filebundle_root = &filebundle_entry_%06d;\n' % idx)