formats.py

   1 import os, sys
   2 from support import PipeThroughCommand, escape, Tmp
   3 import rox
   4
   5 current_command = None
   6
   7 def pipe_through_command(command, src, dst):
   8         global current_command
   9         assert not current_command
  10         current_command = PipeThroughCommand(command, src, dst)
  11         try:
  12                 current_command.wait()
  13         finally:
  14                 current_command = None
  15
  16 operations = []
  17 class Operation:
  18         add_extension = 0
  19
  20         def __init__(self, extension):
  21                 operations.append(self)
  22                 self.extension = extension
  23
  24         def can_handle(self, data):
  25                 return isinstance(data, FileData)
  26
  27         def save_to_stream(self, data, stream):
  28                 pipe_through_command(self.command, data.source, stream)
  29
  30 class Compress(Operation):
  31         "Compress a stream into another stream."
  32         add_extension = 1
  33
  34         def __init__(self, extension, command, type):
  35                 Operation.__init__(self, extension)
  36                 self.command = command
  37                 self.type = type
  38
  39         def __str__(self):
  40                 return 'Compress as .%s' % self.extension
  41
  42 class Decompress(Operation):
  43         "Decompress a stream into another stream."
  44         type = 'text/plain'
  45
  46         def __init__(self, extension, command):
  47                 Operation.__init__(self, extension)
  48                 self.command = command
  49
  50         def __str__(self):
  51                 return 'Decompress .%s' % self.extension
  52
  53 class Extract(Operation):
  54         "Extract an archive to a directory."
  55         type = 'inode/directory'
  56
  57         def __init__(self, extension, command):
  58                 "If command has a %s then the source path is inserted, else uses stdin."
  59                 Operation.__init__(self, extension)
  60                 self.command = command
  61
  62         def __str__(self):
  63                 return 'Extract from a .%s' % self.extension
  64
  65         def save_to_stream(self, data, stream):
  66                 raise Exception('This operation creates a directory, so you have '
  67                                 'to drag to a filer window on the local machine')
  68
  69         def save_to_file(self, data, path):
  70                 if os.path.exists(path):
  71                         if not os.path.isdir(path):
  72                                 raise Exception("'%s' already exists and is not a directory!" %
  73                                                 path)
  74                 if not os.path.exists(path):
  75                         os.mkdir(path)
  76                 os.chdir(path)
  77                 command = self.command
  78                 source = data.source
  79                 if command.find("'%s'") != -1:
  80                         command = command % escape(source.name)
  81                         source = None
  82                 try:
  83                         pipe_through_command(command, source, None)
  84                 finally:
  85                         try:
  86                                 os.rmdir(path)  # Will only succeed if it's empty
  87                         except:
  88                                 pass
  89                 # If we created only a single subdirectory, move it up.
  90                 dirs = os.listdir(path)
  91                 if len(dirs) != 1:
  92                         return
  93                 dir = dirs[0]
  94                 unneeded_path = os.path.join(path, dir)
  95                 if not os.path.isdir(unneeded_path):
  96                         return
  97                 import random
  98                 tmp_path = os.path.join(path, 'tmp-' + `random.randint(0, 100000)`)
  99                 os.rename(unneeded_path, tmp_path)
 100                 for file in os.listdir(tmp_path):
 101                         os.rename(os.path.join(tmp_path, file), os.path.join(path, file))
 102                 os.rmdir(tmp_path)
 103
 104 class Archive(Operation):
 105         "Create an archive from a directory."
 106         add_extension = 1
 107
 108         def __init__(self, extension, command, type):
 109                 assert command.find("'%s'") != -1
 110
 111                 Operation.__init__(self, extension)
 112                 self.command = command
 113                 self.type = type
 114
 115         def __str__(self):
 116                 return 'Create .%s archive' % self.extension
 117
 118         def can_handle(self, data):
 119                 return isinstance(data, DirData)
 120
 121         def save_to_stream(self, data, stream):
 122                 os.chdir(os.path.dirname(data.path))
 123                 command = self.command % escape(os.path.basename(data.path))
 124                 pipe_through_command(command, None, stream)
 125
 126 tgz = Extract('tgz',    "gunzip -c - | tar xf -")
 127 tbz = Extract('tar.bz2', "bunzip2 -c - | tar xf -")
 128 rar = Extract('rar',    "rar x -")
 129 tar = Extract('tar',    "tar xf -")
 130 rpm = Extract('rpm',    "rpm2cpio - | cpio -id --quiet")
 131 cpio = Extract('cpio',  "cpio -id --quiet")
 132 deb = Extract('deb', "ar x '%s'")
 133 zip = Extract('zip', "unzip -q '%s'")
 134 jar = Extract('jar', "unzip -q '%s'")
 135
 136 make_tgz = Archive('tgz', "tar cf - '%s' | gzip", 'application/x-compressed-tar')
 137 Archive('tar.gz',  "tar cf - '%s' | gzip", 'application/x-compressed-tar')
 138 Archive('tar.bz2', "tar cf - '%s' | bzip2", 'application/x-bzip-compressed-tar')
 139 Archive('zip',     "zip -qr - '%s'", 'application/zip'),
 140 Archive('jar',     "zip -qr - '%s'", 'application/x-jar')
 141 Archive('tar',     "tar cf - '%s'", 'application/x-tar')
 142
 143 # Note: these go afterwards so that .tar.gz matches before .gz
 144 make_gz = Compress('gz', "gzip -c -", 'application/x-gzip')
 145 Compress('bz2', "bzip2 -c -", 'application/x-bzip')
 146
 147 gz = Decompress('gz',   "gunzip -c -")
 148 bz2 = Decompress('bz2',  "bunzip2 -ck -")
 149
 150
 151 # Can bzip2 read bzip files?
 152
 153 aliases = {
 154         'tar.gz': 'tgz',
 155         'tar.bz': 'tar.bz2',
 156         'bz': 'bz2'
 157 }
 158
 159 known_extensions = {}
 160 for x in operations:
 161         try:
 162                 known_extensions[x.extension] = None
 163         except AttributeError:
 164                 pass
 165
 166 class FileData:
 167         "A file on the local filesystem."
 168         def __init__(self, path):
 169                 self.path = path
 170
 171                 if path == '-':
 172                         source = sys.stdin
 173                 else:
 174                         try:
 175                                 source = file(path)
 176                         except:
 177                                 rox.report_exception()
 178                                 sys.exit(1)
 179
 180                 self.path = path
 181                 start = source.read(300)
 182                 try:
 183                         source.seek(0)
 184                         self.source = source
 185                 except:
 186                         # Input is not a regular, local, seekable file, so copy it
 187                         # to a local temp file.
 188                         import shutil
 189                         tmp = Tmp()
 190                         tmp.write(start)
 191                         tmp.flush()
 192                         shutil.copyfileobj(source, tmp)
 193                         self.source = tmp
 194                         self.source.seek(0)
 195                 self.default = self.guess_format(start)
 196
 197                 if path == '-':
 198                         name = 'Data'
 199                 else:
 200                         name = path
 201                         for ext in known_extensions:
 202                                 if path.endswith('.' + ext):
 203                                         new = path[:-len(ext)-1]
 204                                         if len(new) < len(name):
 205                                                 name = new
 206                 if self.default.add_extension:
 207                         name += '.' + self.default.extension
 208                 self.default_name = name
 209
 210         def guess_format(self, data):
 211                 "Return a good default Operation, judging by the first 300 bytes or so."
 212                 l = len(data)
 213                 def string(offset, match):
 214                         return data[offset:offset + len(match)] == match
 215                 def short(offset, match):
 216                         if l > offset + 1:
 217                                 a = data[offset]
 218                                 b = data[offset + 1]
 219                                 return ((a == match & 0xff) and (b == (match >> 8))) or \
 220                                         (b == match & 0xff) and (a == (match >> 8))
 221                         return 0
 222
 223                 # Archives
 224                 if string(257, 'ustar\0') or string(257, 'ustar\040\040\0'):
 225                         return tar
 226                 if short(0, 070707) or short(0, 0143561) or string(0, '070707') or \
 227                    string(0, '070701') or string(0, '070702'):
 228                         return cpio
 229                 if string(0, '!<arch>') or string(0, '\\<ar>') or string(0, '<ar>'):
 230                         if string(7, '\ndebian'):
 231                                 return deb
 232                 if string(0, 'Rar!'): return rar
 233                 if string(0, 'PK\003\004'): return zip
 234
 235                 # Compressed streams
 236                 if string(0, '\037\213'):
 237                         if self.path.endswith('.tar.gz') or self.path.endswith('.tgz'):
 238                                 return tgz
 239                         return gz
 240                 if string(0, 'BZh') or string(0, 'BZ'):
 241                         if self.path.endswith('.tar.bz') or self.path.endswith('.tar.bz2') or \
 242                            self.path.endswith('.tbz') or self.path.endswith('.tbz2'):
 243                                 return tbz
 244                         return bz2
 245
 246                 return make_gz
 247
 248 class DirData:
 249         def __init__(self, path):
 250                 self.path = path
 251                 self.default = make_tgz
 252                 self.default_name = path + '.' + self.default.extension