From eb3901a8c96e0c421f3c2232bb0a37646d2b210e Mon Sep 17 00:00:00 2001
From: Christian Muehlhaeuser <muesli@gmail.com>
Date: Thu, 8 Aug 2019 06:18:20 +0200
Subject: [PATCH] Added -search to search for toots containing a particular
 string

---
 main.go | 115 ++++++++++++++++++++++++++++++++++++++++++++------------
 1 file changed, 91 insertions(+), 24 deletions(-)

diff --git a/main.go b/main.go
index d389d7c..e5ee256 100644
--- a/main.go
+++ b/main.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"math"
 	"os"
+	"strings"
 	"time"
 
 	mastodon "github.com/mattn/go-mastodon"
@@ -21,6 +22,7 @@ var (
 	maxToots   = flag.Int("recent", 0, "only account for the N most recent toots (excl replies & boosts)")
 	columns    = flag.Int("columns", 80, "displays tables with N columns")
 	configFile = flag.String("config", "mastodon.json", "uses the specified config file")
+	search     = flag.String("search", "", "searches toots containing string")
 	// user     = flag.String("user", "@fribbledom@mastodon.social", "shows stats for this user")
 )
 
@@ -42,7 +44,7 @@ func registerApp(config *Config) (string, error) {
 	return app.AuthURI, nil
 }
 
-func initClient() {
+func initClient() error {
 	var err error
 	var instance, token, redirectURI, authURI, id, secret string
 	config, err := LoadConfig(*configFile)
@@ -59,7 +61,7 @@ func initClient() {
 		fmt.Print("Which instance to connect to (e.g. https://mastodon.social): ")
 		scanner.Scan()
 		if scanner.Err() != nil {
-			panic(err)
+			return fmt.Errorf("Can't open input: %s", err)
 		}
 		instance = scanner.Text()
 		config.Set("instance", instance)
@@ -68,7 +70,7 @@ func initClient() {
 	if len(id) == 0 {
 		authURI, err = registerApp(&config)
 		if err != nil {
-			panic(err)
+			return fmt.Errorf("Can't register app: %s", err)
 		}
 
 		id = config.Value("id").(string)
@@ -88,37 +90,67 @@ func initClient() {
 		fmt.Printf("Please visit %s and enter the generated token: ", authURI)
 		scanner.Scan()
 		if scanner.Err() != nil {
-			panic(err)
+			return fmt.Errorf("Can't open input: %s", err)
 		}
 		code := scanner.Text()
 
 		err = client.AuthenticateToken(context.Background(), code, redirectURI)
 		if err != nil {
-			panic(err)
+			return fmt.Errorf("Can't retrieve authentication token: %s", err)
 		}
 
 		config.Set("token", mConfig.AccessToken)
-		config.Save(*configFile)
+		err = config.Save(*configFile)
+		if err != nil {
+			return fmt.Errorf("Can't save config: %s", err)
+		}
 	}
+
+	return nil
 }
 
-func main() {
-	flag.Parse()
-	initClient()
-
-	var err error
-	self, err = client.GetAccountCurrentUser(context.Background())
-	if err != nil {
-		panic(err)
+func searchToots() error {
+	pb := &goprogressbar.ProgressBar{
+		Text:  fmt.Sprintf("Searching toots for %s", *search),
+		Total: self.StatusesCount,
+		PrependTextFunc: func(p *goprogressbar.ProgressBar) string {
+			return fmt.Sprintf("%d of %d", p.Current, int64(math.Max(float64(p.Current), float64(self.StatusesCount))))
+		},
+		Current: 0,
+		Width:   40,
 	}
-	/*
-		accounts, err := client.AccountsSearch(context.Background(), *user, 1)
-		if err != nil {
-			panic(err)
-		}
-		self = accounts[0]
-	*/
 
+	var pg mastodon.Pagination
+	for {
+		pg.Limit = 40
+		statuses, err := client.GetAccountStatuses(context.Background(), self.ID, &pg)
+		if err != nil {
+			return fmt.Errorf("Can't retrieve statuses: %s", err)
+		}
+
+		abort := false
+		for _, s := range statuses {
+			if strings.Contains(strings.ToLower(cleanupContent(s.Content)), *search) {
+				fmt.Println("\nFound toot:", cleanupContent(s.Content))
+				fmt.Println()
+			}
+
+			pb.Current += 1
+			pb.LazyPrint()
+		}
+
+		// For some reason, either because it's Pleroma or because I have too few toots,
+		// `pg.MaxID` never equals `""` and we get stuck looping forever. Add a simple
+		// break condition on "no statuses fetched" to avoid the issue.
+		if abort || pg.MaxID == "" || len(statuses) == 0 {
+			break
+		}
+	}
+
+	return nil
+}
+
+func gatherStats() error {
 	stats := &stats{
 		DaysActive: int(time.Since(self.CreatedAt).Hours() / 24),
 		Followers:  self.FollowersCount,
@@ -145,14 +177,14 @@ func main() {
 		pg.Limit = 40
 		statuses, err := client.GetAccountStatuses(context.Background(), self.ID, &pg)
 		if err != nil {
-			panic(err)
+			return fmt.Errorf("Can't retrieve statuses: %s", err)
 		}
 
 		abort := false
 		for _, s := range statuses {
 			err = parseToot(s, stats)
 			if err != nil {
-				panic(err)
+				return fmt.Errorf("Can't parse toot: %s", err)
 			}
 
 			pb.Current += 1
@@ -170,12 +202,47 @@ func main() {
 		if abort || pg.MaxID == "" || len(statuses) == 0 {
 			break
 		}
-		time.Sleep(1000 * time.Millisecond)
 	}
 
+	// print out stats we gathered
 	fmt.Printf("\n\n")
 	printAccountStats(stats)
 	printInteractionStats(stats)
 	printTootStats(stats)
 	printTagStats(stats)
+
+	return nil
+}
+
+func main() {
+	flag.Parse()
+	*search = strings.ToLower(*search)
+	if err := initClient(); err != nil {
+		fmt.Println(err)
+		os.Exit(1)
+	}
+
+	var err error
+	self, err = client.GetAccountCurrentUser(context.Background())
+	if err != nil {
+		fmt.Printf("Can't retrieve user: %s\n", err)
+		os.Exit(1)
+	}
+	/*
+		accounts, err := client.AccountsSearch(context.Background(), *user, 1)
+		if err != nil {
+			panic(err)
+		}
+		self = accounts[0]
+	*/
+
+	if len(*search) > 0 {
+		err = searchToots()
+	} else {
+		err = gatherStats()
+	}
+	if err != nil {
+		fmt.Println(err)
+		os.Exit(1)
+	}
 }