pyfec: add -f option to fec, add more user-friendly handling of filesystem errors and user errors

pyfec/bin/fec

@@ -22,8 +22,9 @@ parser.add_argument('inputfile', help='file to encode or "-" for stdin', type=ar
 parser.add_argument('-d', '--output-dir', help='directory in which share file names will be created (default ".")', default='.', metavar='D')
 parser.add_argument('-p', '--prefix', help='prefix for share file names; If omitted, the name of the input file will be used.', metavar='P')
 parser.add_argument('-s', '--suffix', help='suffix for share file names (default ".fec")', default='.fec', metavar='S')
-parser.add_argument('-m', '--totalshares', help='the total number of share files created (default 16)', default=16, metavar='M')
-parser.add_argument('-k', '--requiredshares', help='the number of share files required to reconstruct (default 4)', default=4, metavar='K')
+parser.add_argument('-m', '--totalshares', help='the total number of share files created (default 16)', default=16, type=int, metavar='M')
+parser.add_argument('-k', '--requiredshares', help='the number of share files required to reconstruct (default 4)', default=4, type=int, metavar='K')
+parser.add_argument('-f', '--force', help='overwrite any file which already in place an output file (share file)', action='store_true')
 parser.add_argument('-v', '--verbose', help='print out messages about progress', action='store_true')
 parser.add_argument('-V', '--version', help='print out version number and exit', action='store_true')
 args = parser.parse_args()
@@ -33,13 +34,22 @@ if args.prefix is None:
     if args.prefix == "<stdin>":
         args.prefix = ""
 
-if args.totalshares < 3 or args.totalshares > 256 or args.requiredshares < 2 or args.requiredshares >= args.totalshares:
-    print "Invalid parameters, requiredshares: %s, totalshares:%s\nPlease see the accompanying documentation." % (args.requiredshares, args.totalshares,)
+if args.totalshares < 3:
+    print "Invalid parameters, totalshares is required to be >= 3\nPlease see the accompanying documentation."
+    sys.exit(1)
+if args.totalshares > 256:
+    print "Invalid parameters, totalshares is required to be <= 256\nPlease see the accompanying documentation."
+    sys.exit(1)
+if args.requiredshares < 2:
+    print "Invalid parameters, requiredshares is required to be >= 2\nPlease see the accompanying documentation."
+    sys.exit(1)
+if args.requiredshares >= args.totalshares:
+    print "Invalid parameters, requiredshares is required to be < totalshares\nPlease see the accompanying documentation."
     sys.exit(1)
 
 args.inputfile.seek(0, 2)
 fsize = args.inputfile.tell()
 args.inputfile.seek(0, 0)
-ret = filefec.encode_to_files(args.inputfile, fsize, args.output_dir, args.prefix, args.requiredshares, args.totalshares, args.suffix, args.verbose)
+ret = filefec.encode_to_files(args.inputfile, fsize, args.output_dir, args.prefix, args.requiredshares, args.totalshares, args.suffix, args.force, args.verbose)
 
 sys.exit(ret)

pyfec/bin/unfec

@@ -40,6 +40,10 @@ else:
         sys.exit(2)
     outf = os.fdopen(outfd, "wb")
 
-ret = filefec.decode_from_files(outf, args.sharefiles, args.verbose)
+try:
+    ret = filefec.decode_from_files(outf, args.sharefiles, args.verbose)
+except filefec.InsufficientShareFilesError, e:
+    print str(e)
+    sys.exit(3)
 
 sys.exit(ret)

pyfec/fec/filefec.py

@@ -31,6 +31,21 @@ import array, os, re, struct, traceback
 
 CHUNKSIZE = 4096
 
+class InsufficientShareFilesError(fec.Error):
+    def __init__(self, k, kb, *args, **kwargs):
+        fec.Error.__init__(self, *args, **kwargs)
+        self.k = k
+        self.kb = kb
+
+    def __repr__(self):
+        return "Insufficient share files -- %d share files are required to recover this file, but only %d were given" % (self.k, self.kb,)
+
+    def __str__(self):
+        return self.__repr__()
+
+class CorruptedShareFilesError(fec.Error):
+    pass
+
 def _build_header(m, k, pad, sh):
     """
     @param m: the total number of shares; 3 <= m <= 256
@@ -106,7 +121,7 @@ def _parse_header(inf):
     # The first 8 bits always encode m.
     ch = inf.read(1)
     if not ch:
-        raise fec.Error("Share files were corrupted -- share file %r didn't have a complete metadata header at the front.  Perhaps the file was truncated." % (inf.name,))
+        raise CorruptedShareFilesError("Share files were corrupted -- share file %r didn't have a complete metadata header at the front.  Perhaps the file was truncated." % (inf.name,))
     byte = ord(ch)
     m = byte + 3
 
@@ -116,7 +131,7 @@ def _parse_header(inf):
     kbitmask = MASK(kbits) << b2_bits_left
     ch = inf.read(1)
     if not ch:
-        raise fec.Error("Share files were corrupted -- share file %r didn't have a complete metadata header at the front.  Perhaps the file was truncated." % (inf.name,))
+        raise CorruptedShareFilesError("Share files were corrupted -- share file %r didn't have a complete metadata header at the front.  Perhaps the file was truncated." % (inf.name,))
     byte = ord(ch)
     k = ((byte & kbitmask) >> b2_bits_left) + 2
 
@@ -129,7 +144,7 @@ def _parse_header(inf):
     if needed_padbits > 0:
         ch = inf.read(1)
         if not ch:
-            raise fec.Error("Share files were corrupted -- share file %r didn't have a complete metadata header at the front.  Perhaps the file was truncated." % (inf.name,))
+            raise CorruptedShareFilesError("Share files were corrupted -- share file %r didn't have a complete metadata header at the front.  Perhaps the file was truncated." % (inf.name,))
         byte = struct.unpack(">B", ch)[0]
         val <<= 8
         val |= byte 
@@ -143,7 +158,7 @@ def _parse_header(inf):
     if needed_shbits > 0:
         ch = inf.read(1)
         if not ch:
-            raise fec.Error("Share files were corrupted -- share file %r didn't have a complete metadata header at the front.  Perhaps the file was truncated." % (inf.name,))
+            raise CorruptedShareFilesError("Share files were corrupted -- share file %r didn't have a complete metadata header at the front.  Perhaps the file was truncated." % (inf.name,))
         byte = struct.unpack(">B", ch)[0]
         val <<= 8
         val |= byte 
@@ -158,7 +173,7 @@ def _parse_header(inf):
 
 FORMAT_FORMAT = "%%s.%%0%dd_%%0%dd%%s"
 RE_FORMAT = "%s.[0-9]+_[0-9]+%s"
-def encode_to_files(inf, fsize, dirname, prefix, k, m, suffix=".fec", verbose=False):
+def encode_to_files(inf, fsize, dirname, prefix, k, m, suffix=".fec", overwrite=False, verbose=False):
     """
     Encode inf, writing the shares to specially named, newly created files.
 
@@ -181,8 +196,11 @@ def encode_to_files(inf, fsize, dirname, prefix, k, m, suffix=".fec", verbose=Fa
             fn = os.path.join(dirname, format % (prefix, shnum, m, suffix,))
             if verbose:
                 print "Creating share file %r..." % (fn,)
-            fd = os.open(fn, os.O_WRONLY|os.O_CREAT|os.O_EXCL)
-            f = os.fdopen(fd, "wb")
+            if overwrite:
+                f = open(fn, "wb")
+            else:
+                fd = os.open(fn, os.O_WRONLY|os.O_CREAT|os.O_EXCL)
+                f = os.fdopen(fd, "wb")
             f.write(hdr)
             fs.append(f)
             fns.append(fn)
@@ -242,13 +260,15 @@ def decode_from_files(outf, infiles, verbose=False):
     for f in infiles:
         (nm, nk, npadlen, shnum,) = _parse_header(f)
         if not (m is None or m == nm):
-            raise fec.Error("Share files were corrupted -- share file %r said that m was %s but another share file previously said that m was %s" % (f.name, nm, m,))
+            raise CorruptedShareFilesError("Share files were corrupted -- share file %r said that m was %s but another share file previously said that m was %s" % (f.name, nm, m,))
         m = nm
         if not (k is None or k == nk):
-            raise fec.Error("Share files were corrupted -- share file %r said that k was %s but another share file previously said that k was %s" % (f.name, nk, k,))
+            raise CorruptedShareFilesError("Share files were corrupted -- share file %r said that k was %s but another share file previously said that k was %s" % (f.name, nk, k,))
+        if k > len(infiles):
+            raise InsufficientShareFilesError(k, len(infiles))
         k = nk
         if not (padlen is None or padlen == npadlen):
-            raise fec.Error("Share files were corrupted -- share file %r said that pad length was %s but another share file previously said that pad length was %s" % (f.name, npadlen, padlen,))
+            raise CorruptedShareFilesError("Share files were corrupted -- share file %r said that pad length was %s but another share file previously said that pad length was %s" % (f.name, npadlen, padlen,))
         padlen = npadlen
 
         infs.append(f)
@@ -262,7 +282,7 @@ def decode_from_files(outf, infiles, verbose=False):
     while True:
         chunks = [ inf.read(CHUNKSIZE) for inf in infs ]
         if [ch for ch in chunks if len(ch) != len(chunks[-1])]:
-            raise fec.Error("Share files were corrupted -- all share files are required to be the same length, but they weren't.")
+            raise CorruptedShareFilesError("Share files were corrupted -- all share files are required to be the same length, but they weren't.")
 
         if len(chunks[-1]) == CHUNKSIZE:
             # Then this was a full read, so we're still in the sharefiles.