#!/usr/bin/env python
'''Extract request records from Common Crawl WARC-format files
given length, offset and file triples.
Input one triple on command line, or
triples from stdin as tab-delimited lines
or complete cdx index lines.

Note that if no output flag(s) is/are given, the whole WARC record will be output, more efficiently than would be the case if -whb is given.'''

import sys, argparse, regex

HACK_USAGE=regex.compile('\[-x\]\n\s*\[length\] \[offset\] \[filename\]')

class HackFormat(argparse.RawDescriptionHelpFormatter):
  def format_help(self):
    global FOO
    FOO=argparse.RawDescriptionHelpFormatter.format_help(self)
    return HACK_USAGE.sub('\n             [ ( -x | length offset filename ) ]',
                          FOO)

def process(options,buf,file,offset,length,whole):
  if whole:
    file.seek(offset)
    bv=memoryview(buf)[:length]
    nb=file.readinto(bv)
    if nb!=length:
      print("losing",file.name,length,nb,file=sys.stderr)
      exit(1)
    sys.stdout.buffer.write(bv)
  file.close()

def main():
  parser = argparse.ArgumentParser(
    description='''Extract records from warc files given length, offset and file triples.
  Input one triple on command line, or
  triples from stdin as tab-delimited lines
  or complete cdx index lines.''',
    epilog='''Note that if no output flag(s) is/are given,
  the whole WARC record will be output, more efficiently than
  would be the case if all three flags were given.''',
    add_help=False,
    conflict_handler='resolve',
    formatter_class=HackFormat
    )

  parser.add_argument('--help',help='Show help',action='help')
  parser.add_argument('-d','--debug',help='Debug output',action='store_true')
  parser.add_argument('-w','--warc',help='output WARC headers',
                      action='store_true')
  parser.add_argument('-h','--headers',help='output HTTP headers',
                      action='store_true')
  parser.add_argument('-b','--body',help='output HTTP body',
                      action='store_true')
  parser.add_argument('-c','--cmd',help='pipes each result thru CMD')
  sg=parser.add_mutually_exclusive_group()
  sg.add_argument('-x','--index',
                      help='take lines of triples from a cdx index file as input',
                      action='store_true')
  sg.add_argument('length',type=int,
                  help='length in bytes of gzipped record',
                  nargs='?')
  parser.add_argument('offset',type=int,
                      help='start position in bytes of gzipped record',
                      nargs='?')
  parser.add_argument('filename',
                      help='name of gzipped Common Crawl WARC-format file',
                      nargs='?',
                      type=argparse.FileType('rb',0))
  # Hack the order of optional and positional in the help output
  parser._action_groups.sort(key=lambda g:g.title)
  #parser.print_help()
  pa=parser.parse_args(sys.argv[1:])
  #print(pa,file=sys.stderr)
  if pa.length is not None:
    # We have to enforce our own check..
    if pa.offset is None or pa.filename is None:
      parser.error("length, offset and filename must all be supplied together")
 
  buf=bytearray(1024*1024)

  whole=not (pa.warc or pa.headers or pa.body)
  if pa.length is not None:
    process(pa,buf,pa.filename,pa.offset,pa.length,whole)
    exit(0)
  if pa.index:
    CDX=regex.compile('length": "([0-9]*)", "offset": "([0-9]*)", "filename": "crawl-data/([^/]*)/segments/([^/]*)/warc/(.*\.gz)"')
    for l in sys.stdin:
      m=CDX.search(l)
      if m is None:
        print("index line problem: \"%s\""%l.lstrip(),file=sys.stderr)
        exit(2)
      f="%s/%s/%s"%(m[3],m[4],m[5])
      process(pa,buf,open(f,'rb',0),int(m[2]),int(m[1]),whole)
    exit(0)
if __name__ == "__main__":
    main()