Multistream WARCs and better Zstandard support

[tofuproxy.git] / cmd / warc-extract / main.go

diff --git a/cmd/warc-extract/main.go b/cmd/warc-extract/main.go
index 8b1aa2ab99db21fcc0afab6ac52f412c9e6375bb..177d251e3ba2967a94388a2894fb1210caa2c739 100644 (file)
--- a/cmd/warc-extract/main.go
+++ b/cmd/warc-extract/main.go
@@ -18,6 +18,9 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 package main
 
 import (
+       "bufio"
+       "bytes"
+       "encoding/binary"
        "flag"
        "fmt"
        "io"
@@ -32,9 +35,53 @@ func main() {
        uri := flag.String("uri", "", "URI to extract, if specified")
        hdr := flag.Bool("hdr", false, "Also extract WARC's header")
        idx := flag.Bool("idx", false, "Save WARC indexes")
+       recompress := flag.Bool("for-enzstd", false, "Output for enzstd utility")
        flag.Parse()
        log.SetFlags(log.Lshortfile)
 
+       if *recompress {
+               var hdr bytes.Buffer
+               size := make([]byte, 8)
+               bw := bufio.NewWriter(os.Stdout)
+               for _, p := range flag.Args() {
+                       r, err := warc.NewReader(p)
+                       if err != nil {
+                               log.Fatalln(err)
+                       }
+                       for {
+                               rec, rr, err := r.ReadRecord()
+                               if err != nil {
+                                       if err == io.EOF {
+                                               break
+                                       }
+                                       log.Fatalln(err)
+                               }
+                               for _, line := range rec.HdrLines {
+                                       hdr.WriteString(line)
+                               }
+                               hdr.WriteString("\r\n")
+                               binary.BigEndian.PutUint64(size, uint64(hdr.Len())+uint64(rec.Size)+4)
+                               if _, err = bw.Write(size); err != nil {
+                                       log.Fatalln(err)
+                               }
+                               if _, err = io.Copy(bw, &hdr); err != nil {
+                                       log.Fatalln(err)
+                               }
+                               if _, err = io.Copy(bw, rr); err != nil {
+                                       log.Fatalln(err)
+                               }
+                               r.RecordWasRead()
+                               if _, err = bw.Write([]byte("\r\n\r\n")); err != nil {
+                                       log.Fatalln(err)
+                               }
+                       }
+               }
+               if err := bw.Flush(); err != nil {
+                       log.Fatalln(err)
+               }
+               return
+       }
+
        for _, p := range flag.Args() {
                log.Println("adding", p)
                if err := warc.Add(p); err != nil {
@@ -63,7 +110,7 @@ func main() {
                        if rec == nil {
                                continue
                        }
-                       r, err := rec.Reader(!*hdr)
+                       r, err := rec.Reader(!*hdr, warc.WARCsOffsets)
                        if err != nil {
                                log.Fatalln(err)
                        }