From 55f90de6617499e7c58e10902159db73efdae0b0 Mon Sep 17 00:00:00 2001
From: Farhan Khan <khanzf@gmail.com>
Date: Wed, 9 Feb 2022 23:09:13 -0500
Subject: [PATCH] Adding Pleroma websocket streaming

---
 fedilogue/instance.go |  3 +-
 fedilogue/stream.go   | 65 +++++++++++++++++++++++++++++++++++++++++--
 go.mod                |  2 ++
 go.sum                |  2 ++
 shared/headers.go     |  5 ++++
 5 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/fedilogue/instance.go b/fedilogue/instance.go
index abe2ca5..63319ec 100644
--- a/fedilogue/instance.go
+++ b/fedilogue/instance.go
@@ -250,7 +250,8 @@ func StartInstance(endpoint string) {
 		logConn("Starting " + endpoint + " as " + o.Software + " " + o.Version)
 		o.CaptureType = "Poll"
 		UpdateRunner(endpoint, o)
-		PollMastodonPleroma(endpoint, &o)
+		//		PollMastodonPleroma(endpoint, &o)
+		StreamPleroma(endpoint)
 	} else if o.Software == "mastodon" {
 		logConn("Starting " + endpoint + " as " + o.Software + " " + o.Version)
 		o.CaptureType = "Stream"
diff --git a/fedilogue/stream.go b/fedilogue/stream.go
index 7dd5bb6..d2a5850 100644
--- a/fedilogue/stream.go
+++ b/fedilogue/stream.go
@@ -3,22 +3,81 @@ package main
 import (
 	"bufio"
 	"encoding/json"
+	"fmt"
+	"log"
+	// Must be removed later
 	"net/http"
 	"strings"
 	"time"
 
+	"github.com/gorilla/websocket"
 	"gitlab.com/khanzf/fedilogue/shared"
 )
 
+func StreamPleroma(endpoint string) {
+	wss_url := "wss://" + endpoint + "/api/v1/streaming/?stream=public"
+	var pleromaHeader shared.PleromaStreamHeader
+	var newactivity shared.ReportActivity
+	var err error
+
+	for {
+		var tries int
+		var ws *websocket.Conn
+		for tries = 0; tries < 10; tries++ {
+			ws, _, err = websocket.DefaultDialer.Dial(wss_url, nil)
+			if err != nil {
+				continue
+			}
+			break
+		}
+
+		if tries == 10 {
+			logWarn("Unable to connect to " + endpoint + " after 10 tries, exiting")
+			return
+		}
+
+		for {
+			_, p, err := ws.ReadMessage()
+			if err != nil {
+				log.Fatal(err)
+				break
+			}
+
+			err = json.Unmarshal(p, &pleromaHeader)
+			if err != nil {
+				logDebug("Unable to parse data from " + endpoint + ", but still connected.")
+				break
+			}
+
+			switch pleromaHeader.Event {
+			case "update":
+				err = json.Unmarshal([]byte(pleromaHeader.Payload), &newactivity)
+				if err != nil {
+					logDebug("Unable to parse data from " + endpoint + ", but still connected.")
+					break
+				}
+				go check_activity(newactivity.Uri)
+			default:
+				continue
+			}
+
+		}
+
+		ws.Close()
+		time.Sleep(time.Minute * 10)
+
+	}
+}
+
 func StreamMastodon(endpoint string, o *shared.RunningInstance) {
 	stream_client := BuildClient(endpoint)
 
 	var oauthData OAuth
 	var retry bool
 
-	for {
+	api_timeline := "https://" + endpoint + "/api/v1/streaming/public"
 
-		api_timeline := "https://" + endpoint + "/api/v1/streaming/public"
+	for {
 		req, err := http.NewRequest("GET", api_timeline, nil)
 		req.Header.Set("User-Agent", "Tusky")
 		if err != nil {
@@ -113,7 +172,7 @@ func StreamMastodon(endpoint string, o *shared.RunningInstance) {
 			}
 		}
 		if retry == true {
-			time.Sleep(time.Minute * 30)
+			time.Sleep(time.Minute * 10)
 		} else {
 			break
 		}
diff --git a/go.mod b/go.mod
index 9015141..10f4808 100644
--- a/go.mod
+++ b/go.mod
@@ -9,6 +9,8 @@ require (
 	muzzammil.xyz/jsonc v0.0.0-20201229145248-615b0916ca38
 )
 
+require github.com/gorilla/websocket v1.4.2
+
 require (
 	github.com/aymerick/douceur v0.2.0 // indirect
 	github.com/gorilla/css v1.0.0 // indirect
diff --git a/go.sum b/go.sum
index 5a10a7a..2019d6a 100644
--- a/go.sum
+++ b/go.sum
@@ -20,6 +20,8 @@ github.com/gofrs/uuid v4.0.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRx
 github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
 github.com/gorilla/css v1.0.0 h1:BQqNyPTi50JCFMTw/b67hByjMVXZRwGha6wxVGkeihY=
 github.com/gorilla/css v1.0.0/go.mod h1:Dn721qIggHpt4+EFCcTLTU/vk5ySda2ReITrtgBl60c=
+github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
+github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/jackc/chunkreader v1.0.0 h1:4s39bBR8ByfqH+DKm8rQA3E1LHZWB9XWcrz8fqaZbe0=
 github.com/jackc/chunkreader v1.0.0/go.mod h1:RT6O25fNZIuasFJRyZ4R/Y2BbhasbmZXF9QQ7T3kePo=
 github.com/jackc/chunkreader/v2 v2.0.0/go.mod h1:odVSm741yZoC3dpHEUXIqA9tQRhFrgOHwnPIn9lDKlk=
diff --git a/shared/headers.go b/shared/headers.go
index 67f8375..305c463 100644
--- a/shared/headers.go
+++ b/shared/headers.go
@@ -40,6 +40,11 @@ type ReportActivity struct {
 	normalized string
 }
 
+type PleromaStreamHeader struct {
+	Event   string `json:"event"`
+	Payload string `json:"payload"`
+}
+
 type AccountType struct {
 	Acct         string `json:"acct"`
 	Avatar       string `json:"avatar"`