From ca6c24dee10e26d3e29ea7d38ae5da25da37dfdf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Sat, 17 Sep 2022 18:17:04 +0200
Subject: [PATCH] add screenshot function.

---
 docs/changelog.md                      |  1 +
 server/internal/http/http.go           | 37 +++++++++++++++++++++++++-
 server/internal/types/websocket.go     |  1 +
 server/internal/websocket/websocket.go |  4 +++
 server/neko.go                         |  2 +-
 5 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/docs/changelog.md b/docs/changelog.md
index a08c1e1..f3678b0 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -6,6 +6,7 @@
 - Added `m1k1o/neko:vivaldi` tag (thanks @Xeddius).
 - Added `m1k1o/neko:opera` tag (thanks @prophetofxenu).
 - Added `NEKO_PATH_PREFIX`.
+- Added screenshot function `/screenshot.jpg?pwd=<admin>`, works only for unlocked rooms.
 
 ### Misc
 - Server: Split `remote` to `desktop` and `capture`.
diff --git a/server/internal/http/http.go b/server/internal/http/http.go
index 21c9617..06b0863 100644
--- a/server/internal/http/http.go
+++ b/server/internal/http/http.go
@@ -3,8 +3,10 @@ package http
 import (
 	"context"
 	"encoding/json"
+	"image/jpeg"
 	"net/http"
 	"os"
+	"strconv"
 
 	"github.com/go-chi/chi"
 	"github.com/go-chi/chi/middleware"
@@ -22,7 +24,7 @@ type Server struct {
 	conf   *config.Server
 }
 
-func New(conf *config.Server, webSocketHandler types.WebSocketHandler) *Server {
+func New(conf *config.Server, webSocketHandler types.WebSocketHandler, desktop types.DesktopManager) *Server {
 	logger := log.With().Str("module", "http").Logger()
 
 	router := chi.NewRouter()
@@ -64,6 +66,39 @@ func New(conf *config.Server, webSocketHandler types.WebSocketHandler) *Server {
 		}
 	})
 
+	router.Get("/screenshot.jpg", func(w http.ResponseWriter, r *http.Request) {
+		password := r.URL.Query().Get("pwd")
+		isAdmin, err := webSocketHandler.IsAdmin(password)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusForbidden)
+			return
+		}
+
+		if !isAdmin {
+			http.Error(w, "bad authorization", http.StatusUnauthorized)
+			return
+		}
+
+		if webSocketHandler.IsLocked("login") {
+			http.Error(w, "room is locked", http.StatusLocked)
+			return
+		}
+
+		quality, err := strconv.Atoi(r.URL.Query().Get("quality"))
+		if err != nil {
+			quality = 90
+		}
+
+		w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
+		w.Header().Set("Content-Type", "image/jpeg")
+
+		img := desktop.GetScreenshotImage()
+		if err := jpeg.Encode(w, img, &jpeg.Options{Quality: quality}); err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+	})
+
 	router.Get("/health", func(w http.ResponseWriter, r *http.Request) {
 		_, _ = w.Write([]byte("true"))
 	})
diff --git a/server/internal/types/websocket.go b/server/internal/types/websocket.go
index 0ea4ee6..968bbb8 100644
--- a/server/internal/types/websocket.go
+++ b/server/internal/types/websocket.go
@@ -32,5 +32,6 @@ type WebSocketHandler interface {
 	Shutdown() error
 	Upgrade(w http.ResponseWriter, r *http.Request) error
 	Stats() Stats
+	IsLocked(resource string) bool
 	IsAdmin(password string) (bool, error)
 }
diff --git a/server/internal/websocket/websocket.go b/server/internal/websocket/websocket.go
index 8ec8e62..a4ec572 100644
--- a/server/internal/websocket/websocket.go
+++ b/server/internal/websocket/websocket.go
@@ -298,6 +298,10 @@ func (ws *WebSocketHandler) Stats() types.Stats {
 	}
 }
 
+func (ws *WebSocketHandler) IsLocked(resource string) bool {
+	return ws.state.IsLocked(resource)
+}
+
 func (ws *WebSocketHandler) IsAdmin(password string) (bool, error) {
 	if password == ws.conf.AdminPassword {
 		return true, nil
diff --git a/server/neko.go b/server/neko.go
index 981584e..abaec5f 100644
--- a/server/neko.go
+++ b/server/neko.go
@@ -134,7 +134,7 @@ func (neko *Neko) Start() {
 	webSocketHandler := websocket.New(sessionManager, desktopManager, captureManager, webRTCManager, neko.WebSocket)
 	webSocketHandler.Start()
 
-	server := http.New(neko.Server, webSocketHandler)
+	server := http.New(neko.Server, webSocketHandler, desktopManager)
 	server.Start()
 
 	neko.sessionManager = sessionManager