From d091d4a8bb4d777ac11e1d66712e47f45fd08371 Mon Sep 17 00:00:00 2001
From: Nick Craig-Wood <nick@craig-wood.com>
Date: Wed, 8 Feb 2017 08:09:41 +0000
Subject: [PATCH] rclone cat: add --head, --tail, --offset, --count and
 --discard

Fixes #819
---
 cmd/cat/cat.go        | 55 ++++++++++++++++++++++++++++++++++++++++++-
 fs/operations.go      | 31 ++++++++++++++++++++----
 fs/operations_test.go | 28 +++++++++++++++-------
 3 files changed, 100 insertions(+), 14 deletions(-)

diff --git a/cmd/cat/cat.go b/cmd/cat/cat.go
index d650a59cd..a77e447bf 100644
--- a/cmd/cat/cat.go
+++ b/cmd/cat/cat.go
@@ -1,6 +1,9 @@
 package cat
 
 import (
+	"io"
+	"io/ioutil"
+	"log"
 	"os"
 
 	"github.com/ncw/rclone/cmd"
@@ -8,8 +11,22 @@ import (
 	"github.com/spf13/cobra"
 )
 
+// Globals
+var (
+	head    = int64(0)
+	tail    = int64(0)
+	offset  = int64(0)
+	count   = int64(-1)
+	discard = false
+)
+
 func init() {
 	cmd.Root.AddCommand(commandDefintion)
+	commandDefintion.Flags().Int64VarP(&head, "head", "", head, "Only print the first N characters.")
+	commandDefintion.Flags().Int64VarP(&tail, "tail", "", tail, "Only print the last N characters.")
+	commandDefintion.Flags().Int64VarP(&offset, "offset", "", offset, "Start printing at offset N (or from end if -ve).")
+	commandDefintion.Flags().Int64VarP(&count, "count", "", count, "Only print N characters.")
+	commandDefintion.Flags().BoolVarP(&discard, "discard", "", discard, "Discard the output instead of printing.")
 }
 
 var commandDefintion = &cobra.Command{
@@ -29,12 +46,48 @@ Or like this to output any file in dir or subdirectories.
 Or like this to output any .txt files in dir or subdirectories.
 
     rclone --include "*.txt" cat remote:path/to/dir
+
+Use the --head flag to print characters only at the start, --tail for
+the end and --offset and --count to print a section in the middle.
+Note that if offset is negative it will count from the end, so
+--offset -1 --count 1 is equivalent to --tail 1.
 `,
 	Run: func(command *cobra.Command, args []string) {
+		usedOffset := offset != 0 || count >= 0
+		usedHead := head > 0
+		usedTail := tail > 0
+		if usedHead && usedTail || usedHead && usedOffset || usedTail && usedOffset {
+			log.Fatalf("Can only use one of  --head, --tail or --offset with --count")
+		}
+		if head > 0 {
+			offset = 0
+			count = head
+		}
+		if tail > 0 {
+			offset = -tail
+			count = -1
+		}
 		cmd.CheckArgs(1, 1, command, args)
 		fsrc := cmd.NewFsSrc(args)
+		var w io.Writer = os.Stdout
+		if discard {
+			w = ioutil.Discard
+		}
 		cmd.Run(false, false, command, func() error {
-			return fs.Cat(fsrc, os.Stdout)
+			return fs.Cat(fsrc, w, offset, count)
 		})
 	},
 }
+
+/*
+Try removing buffering to stop the transfer!!!
+
+
+Transferred:   2.847 GBytes (2.555 MBytes/s)
+Errors:                 3
+Checks:                 0
+Transferred:         1844
+Elapsed time:     19m0.8s
+Transferring:
+ *          2001test/rogers-wedding/0016_1~1.jpg: 74% done, 0 Bytes/s, ETA: -
+*/
diff --git a/fs/operations.go b/fs/operations.go
index f95e9d594..0b2423d45 100644
--- a/fs/operations.go
+++ b/fs/operations.go
@@ -5,6 +5,7 @@ package fs
 import (
 	"fmt"
 	"io"
+	"io/ioutil"
 	"log"
 	"mime"
 	"path"
@@ -1236,7 +1237,14 @@ func CleanUp(f Fs) error {
 }
 
 // Cat any files to the io.Writer
-func Cat(f Fs, w io.Writer) error {
+//
+// if offset == 0 it will be ignored
+// if offset > 0 then the file will be seeked to that offset
+// if offset < 0 then the file will be seeked that far from the end
+//
+// if count < 0 then it will be ignored
+// if count >= 0 then only that many characters will be output
+func Cat(f Fs, w io.Writer, offset, count int64) error {
 	var mu sync.Mutex
 	return ListFn(f, func(o Object) {
 		var err error
@@ -1244,14 +1252,24 @@ func Cat(f Fs, w io.Writer) error {
 		defer func() {
 			Stats.DoneTransferring(o.Remote(), err == nil)
 		}()
-		mu.Lock()
-		defer mu.Unlock()
-		in, err := o.Open()
+		thisOffset := offset
+		if thisOffset < 0 {
+			thisOffset += o.Size()
+		}
+		var options []OpenOption
+		if thisOffset > 0 {
+			options = append(options, &SeekOption{Offset: thisOffset})
+		}
+		in, err := o.Open(options...)
 		if err != nil {
 			Stats.Error()
 			ErrorLog(o, "Failed to open: %v", err)
 			return
 		}
+		reader := in
+		if count >= 0 {
+			reader = ioutil.NopCloser(&io.LimitedReader{R: in, N: count})
+		}
 		defer func() {
 			err = in.Close()
 			if err != nil {
@@ -1259,7 +1277,10 @@ func Cat(f Fs, w io.Writer) error {
 				ErrorLog(o, "Failed to close: %v", err)
 			}
 		}()
-		inAccounted := NewAccountWithBuffer(in, o) // account and buffer the transfer
+		inAccounted := NewAccountWithBuffer(reader, o) // account and buffer the transfer
+		// take the lock just before we output stuff, so at the last possible moment
+		mu.Lock()
+		defer mu.Unlock()
 		_, err = io.Copy(w, inAccounted)
 		if err != nil {
 			Stats.Error()
diff --git a/fs/operations_test.go b/fs/operations_test.go
index 43710f81c..52bb7e1ee 100644
--- a/fs/operations_test.go
+++ b/fs/operations_test.go
@@ -667,18 +667,30 @@ func TestDeduplicateRename(t *testing.T) {
 func TestCat(t *testing.T) {
 	r := NewRun(t)
 	defer r.Finalise()
-	file1 := r.WriteBoth("file1", "aaa", t1)
-	file2 := r.WriteBoth("file2", "bbb", t2)
+	file1 := r.WriteBoth("file1", "ABCDEFGHIJ", t1)
+	file2 := r.WriteBoth("file2", "012345678", t2)
 
 	fstest.CheckItems(t, r.fremote, file1, file2)
 
-	var buf bytes.Buffer
-	err := fs.Cat(r.fremote, &buf)
-	require.NoError(t, err)
-	res := buf.String()
+	for _, test := range []struct {
+		offset int64
+		count  int64
+		a      string
+		b      string
+	}{
+		{0, -1, "ABCDEFGHIJ", "012345678"},
+		{0, 5, "ABCDE", "01234"},
+		{-3, -1, "HIJ", "678"},
+		{1, 3, "BCD", "123"},
+	} {
+		var buf bytes.Buffer
+		err := fs.Cat(r.fremote, &buf, test.offset, test.count)
+		require.NoError(t, err)
+		res := buf.String()
 
-	if res != "aaabbb" && res != "bbbaaa" {
-		t.Errorf("Incorrect output from Cat: %q", res)
+		if res != test.a+test.b && res != test.b+test.a {
+			t.Errorf("Incorrect output from Cat(%d,%d): %q", test.offset, test.count, res)
+		}
 	}
 }