package main import ( "bytes" "context" "encoding/base64" "encoding/json" "fmt" "io" "log" "net/http" "path/filepath" "screenmark/screenmark/.gen/haystack/haystack/model" "screenmark/screenmark/agents/client" "screenmark/screenmark/models" "github.com/go-chi/chi/v5" "github.com/go-chi/chi/v5/middleware" "github.com/google/uuid" "github.com/joho/godotenv" ) type TestAiClient struct { ImageInfo client.ImageMessageContent } func (client TestAiClient) GetImageInfo(imageName string, imageData []byte) (client.ImageMessageContent, error) { return client.ImageInfo, nil } func main() { err := godotenv.Load() if err != nil { panic(err) } db, err := models.InitDatabase() if err != nil { panic(err) } imageModel := models.NewImageModel(db) userModel := models.NewUserModel(db) mail, err := CreateMailClient() if err != nil { panic(err) } auth := CreateAuth(mail) eventManager := NewEventManager() go ListenNewImageEvents(db, &eventManager) go ListenProcessingImageStatus(db, &eventManager) r := chi.NewRouter() r.Use(middleware.Logger) r.Use(CorsMiddleware) r.Use(func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Add("Content-Type", "application/json") next.ServeHTTP(w, r) }) }) r.Options("/*", func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }) r.Group(func(r chi.Router) { r.Use(ProtectedRoute) r.Get("/image", func(w http.ResponseWriter, r *http.Request) { userId := r.Context().Value(USER_ID).(uuid.UUID) if err != nil { w.WriteHeader(http.StatusForbidden) fmt.Fprintf(w, "You cannot read this") return } images, err := userModel.ListWithProperties(r.Context(), userId) if err != nil { log.Println(err) w.WriteHeader(http.StatusNotFound) fmt.Fprintf(w, "Something went wrong") return } jsonImages, err := json.Marshal(models.GetTypedImageProperties(images)) if err != nil { log.Println(err) w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "Could not create JSON response for this image") return } w.Write(jsonImages) }) r.Get("/image-properties/{id}", func(w http.ResponseWriter, r *http.Request) { userId := r.Context().Value(USER_ID).(uuid.UUID) stringImageId := r.PathValue("id") imageId, err := uuid.Parse(stringImageId) if err != nil { w.WriteHeader(http.StatusForbidden) fmt.Fprintf(w, "You cannot read this") return } image, err := userModel.ListImageWithProperties(r.Context(), userId, imageId) if err != nil { log.Println(err) w.WriteHeader(http.StatusNotFound) fmt.Fprintf(w, "Something went wrong") return } jsonImages, err := json.Marshal(models.GetTypedImageProperties([]models.ImageWithProperties{image})[0]) if err != nil { log.Println(err) w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "Could not create JSON response for this image") return } w.Write(jsonImages) }) r.Get("/image/{id}", func(w http.ResponseWriter, r *http.Request) { stringImageId := r.PathValue("id") userId := r.Context().Value(USER_ID).(uuid.UUID) imageId, err := uuid.Parse(stringImageId) if err != nil { w.WriteHeader(http.StatusForbidden) fmt.Fprintf(w, "You cannot read this") return } if authorized := imageModel.IsUserAuthorized(r.Context(), imageId, userId); !authorized { w.WriteHeader(http.StatusForbidden) fmt.Fprintf(w, "You cannot read this") return } // TODO: really need authorization here! image, err := imageModel.Get(r.Context(), imageId) if err != nil { log.Println(err) w.WriteHeader(http.StatusNotFound) fmt.Fprintf(w, "Could not get image") return } // TODO: this could be part of the db table extension := filepath.Ext(image.Image.ImageName) extension = extension[1:] w.Header().Add("Content-Type", "image/"+extension) w.Write(image.Image.Image) }) r.Post("/image/{name}", func(w http.ResponseWriter, r *http.Request) { imageName := r.PathValue("name") userId := r.Context().Value(USER_ID).(uuid.UUID) if len(imageName) == 0 { w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "You need to provide a name in the path") return } contentType := r.Header.Get("Content-Type") // TODO: length checks on body // TODO: extract this shit out image := make([]byte, 0) if contentType == "application/base64" { decoder := base64.NewDecoder(base64.StdEncoding, r.Body) buf := &bytes.Buffer{} decodedIamge, err := io.Copy(buf, decoder) if err != nil { log.Println(err) w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "bruh, base64 aint decoding") return } fmt.Println(string(image)) fmt.Println(decodedIamge) image = buf.Bytes() } else if contentType == "application/oclet-stream" { bodyData, err := io.ReadAll(r.Body) if err != nil { log.Println(err) w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "bruh, binary aint binaring") return } // TODO: check headers image = bodyData } else { log.Println("bad stuff?") w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "Bruh, you need oclet stream or base64") return } if err != nil { log.Println("First case") w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "Couldnt read the image from the request body") return } userImage, err := imageModel.Process(r.Context(), userId, model.Image{ Image: image, ImageName: imageName, }) if err != nil { log.Println("Second case") log.Println(err) w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "Could not save image to DB") return } jsonUserImage, err := json.Marshal(userImage) if err != nil { log.Println("Third case") log.Println(err) w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "Could not create JSON response for this image") return } w.WriteHeader(http.StatusCreated) fmt.Fprint(w, string(jsonUserImage)) w.Header().Add("Content-Type", "application/json") }) }) r.Get("/image-events/{id}", func(w http.ResponseWriter, r *http.Request) { // TODO: authentication :) id := r.PathValue("id") imageNotifier, exists := eventManager.listeners[id] if !exists { fmt.Println("Not found!") w.WriteHeader(http.StatusNotFound) return } w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") w.(http.Flusher).Flush() // TODO: get the current status of the image and send it across. ctx, cancel := context.WithCancel(r.Context()) for { select { case <-ctx.Done(): fmt.Fprint(w, "event: close\ndata: Connection closed\n\n") w.(http.Flusher).Flush() cancel() return case data := <-imageNotifier: if data == "" { cancel() continue } fmt.Printf("Status received: %s\n", data) fmt.Fprintf(w, "event: data\ndata: %s\n\n", data) w.(http.Flusher).Flush() if data == "complete" { cancel() } } } }) r.Post("/login", func(w http.ResponseWriter, r *http.Request) { type LoginBody struct { Email string `json:"email"` } loginBody := LoginBody{} err := json.NewDecoder(r.Body).Decode(&loginBody) if err != nil { log.Println(err) w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "Request body was not correct") return } // TODO: validate it's an email auth.CreateCode(loginBody.Email) w.WriteHeader(http.StatusOK) }) r.Post("/code", func(w http.ResponseWriter, r *http.Request) { type CodeBody struct { Email string `json:"email"` Code string `json:"code"` } type CodeReturn struct { Access string `json:"access"` Refresh string `json:"refresh"` } codeBody := CodeBody{} if err := json.NewDecoder(r.Body).Decode(&codeBody); err != nil { log.Println(err) w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "Request body was not correct") return } if err := auth.UseCode(codeBody.Email, codeBody.Code); err != nil { log.Println(err) w.WriteHeader(http.StatusBadRequest) fmt.Fprintf(w, "email or code are incorrect") return } if exists := userModel.DoesUserExist(r.Context(), codeBody.Email); !exists { userModel.Save(r.Context(), model.Users{ Email: codeBody.Email, }) } uuid, err := userModel.GetUserIdFromEmail(r.Context(), codeBody.Email) if err != nil { log.Println(err) w.WriteHeader(http.StatusInternalServerError) fmt.Fprintf(w, "Something went wrong.") return } refresh := CreateRefreshToken(uuid) access := CreateAccessToken(uuid) codeReturn := CodeReturn{ Access: access, Refresh: refresh, } fmt.Println(codeReturn) json, err := json.Marshal(codeReturn) if err != nil { log.Println(err) w.WriteHeader(http.StatusInternalServerError) fmt.Fprintf(w, "Something went wrong.") return } w.WriteHeader(http.StatusOK) w.Header().Add("Content-Type", "application/json") fmt.Fprint(w, string(json)) }) logWriter := DatabaseWriter{ dbPool: db, } r.Route("/logs", createLogHandler(&logWriter)) log.Println("Listening and serving on port 3040.") if err := http.ListenAndServe(":3040", r); err != nil { log.Println(err) return } }