package main import ( "bytes" "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) notifier := NewNotifier[Notification](10) go ListenNewImageEvents(db, ¬ifier) go ListenProcessingImageStatus(db, imageModel, ¬ifier) r := chi.NewRouter() r.Use(middleware.Logger) r.Use(CorsMiddleware) r.Options("/*", func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }) // Temporarily not in protect route because we aren't using cookies. // Therefore they don't get automatically attached to the request. // So cannot send the tokensend the token 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 // } 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.ImageName) if len(extension) == 0 { // Same hack extension = "png" } extension = extension[1:] w.Header().Add("Content-Type", "image/"+extension) w.Write(image.Image) }) r.Group(func(r chi.Router) { r.Use(ProtectedRoute) 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.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 } imageProperties, err := userModel.ListWithProperties(r.Context(), userId) if err != nil { log.Println(err) w.WriteHeader(http.StatusNotFound) fmt.Fprintf(w, "Something went wrong") return } images, err := userModel.GetUserImages(r.Context(), userId) if err != nil { log.Println(err) w.WriteHeader(http.StatusNotFound) fmt.Fprintf(w, "Something went wrong") return } type ImagesReturn struct { UserImages []models.UserImageWithImage ImageProperties []models.TypedProperties } imagesReturn := ImagesReturn{ UserImages: images, ImageProperties: models.GetTypedImageProperties(imageProperties), } jsonImages, err := json.Marshal(imagesReturn) 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.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") fmt.Printf("Content-Type: %s\n", contentType) // 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" || contentType == "image/png" { 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.Route("/notifications", func(r chi.Router) { r.Use(GetUserIdFromUrl) r.Get("/", CreateEventsHandler(¬ifier)) }) 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) }) type CodeReturn struct { Access string `json:"access"` Refresh string `json:"refresh"` } r.Post("/code", func(w http.ResponseWriter, r *http.Request) { type CodeBody struct { Email string `json:"email"` Code string `json:"code"` } 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)) }) r.Get("/demo-login", func(w http.ResponseWriter, r *http.Request) { uuid, err := userModel.GetUserIdFromEmail(r.Context(), "demo@email.com") 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 } }