From eb7a03f971966680b86321bc6d165a423e578a76 Mon Sep 17 00:00:00 2001
From: bignutty <3515180-bignutty@users.noreply.gitlab.com>
Date: Thu, 23 May 2024 22:58:33 +0200
Subject: [PATCH] mono2: screenshot

---
 commands/interaction/slash/utils/screenshot.js | 16 +++++++---------
 commands/message/utils/screenshot.js           | 15 ++++++---------
 labscore/api/obelisk/endpoints.js              |  4 +++-
 labscore/api/obelisk/index.js                  |  7 +++++++
 4 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/commands/interaction/slash/utils/screenshot.js b/commands/interaction/slash/utils/screenshot.js
index 9aa1b0b..56155c2 100644
--- a/commands/interaction/slash/utils/screenshot.js
+++ b/commands/interaction/slash/utils/screenshot.js
@@ -1,7 +1,7 @@
 const { createEmbed } = require("../../../../labscore/utils/embed");
 const { editOrReply } = require("../../../../labscore/utils/message");
 
-const { webshot } = require("../../../../labscore/api/obelisk");
+const { WebUtilsWebPageScreenshot} = require("../../../../labscore/api/obelisk");
 
 const { ApplicationCommandOptionTypes, InteractionCallbackTypes } = require('detritus-client/lib/constants');
 
@@ -32,28 +32,26 @@ module.exports = {
     try {
       const t = Date.now();
 
-      let ss = await webshot(context, args.url, false) // nsfw sites are always blocked
+      let ss = await WebUtilsWebPageScreenshot(context, args.url, false) // nsfw sites are always blocked
 
-      if (ss.response.body.status && ss.response.body.status !== 3) {
-        if (ss.response.body.image) return await editOrReply(context,
+      if (ss.response.body.error) return await editOrReply(context,
           createEmbed("image", context, {
-            url: ss.response.body.image,
+            url: ss.response.body.error.image_url,
             time: ((Date.now() - t) / 1000).toFixed(2)
           })
         )
-        return await editOrReply(context, createEmbed("error", context, "Unable to create screenshot."))
-      }
+
 
       return await editOrReply(context, {
         embeds: [createEmbed("image", context, {
           url: "screenshot.png",
           time: ((Date.now() - t) / 1000).toFixed(2)
         })],
-        files: [{ filename: "screenshot.png", value: ss.response.body }]
+        files: [{ filename: "screenshot.png", value: Buffer.from(ss.response.body.image, 'base64') }]
       })
     } catch (e) {
       console.log(e)
-      return await editOrReply(context, createEmbed("image", context, {
+      return editOrReply(context, createEmbed("image", context, {
         url: "https://bignutty.gitlab.io/webstorage4/v2/assets/screenshot/brand-update-2024/scr_unavailable.png"
       }))
     }
diff --git a/commands/message/utils/screenshot.js b/commands/message/utils/screenshot.js
index 050dd99..dab1146 100644
--- a/commands/message/utils/screenshot.js
+++ b/commands/message/utils/screenshot.js
@@ -4,7 +4,7 @@ const { editOrReply } = require("../../../labscore/utils/message");
 const superagent = require('superagent')
 
 const { Permissions } = require("detritus-client/lib/constants");
-const { webshot } = require("../../../labscore/api/obelisk");
+const { WebUtilsWebPageScreenshot} = require("../../../labscore/api/obelisk");
 
 async function processJob(jobUrl) {
   let job = await superagent.get(jobUrl)
@@ -33,24 +33,21 @@ module.exports = {
     try {
       const t = Date.now();
 
-      let ss = await webshot(context, args.url, context.channel.nsfw)
+      let ss = await WebUtilsWebPageScreenshot(context, args.url, false)
 
-      if (ss.response.body.status && ss.response.body.status !== 3) {
-        if (ss.response.body.image) return await editOrReply(context,
+      if (ss.response.body.error) return await editOrReply(context,
           createEmbed("image", context, {
-            url: ss.response.body.image,
+            url: ss.response.body.error.image_url,
             time: ((Date.now() - t) / 1000).toFixed(2)
           })
-        )
-        return await editOrReply(context, createEmbed("error", context, "Unable to create screenshot."))
-      }
+      )
 
       return await editOrReply(context, {
         embeds: [createEmbed("image", context, {
           url: "screenshot.png",
           time: ((Date.now() - t) / 1000).toFixed(2)
         })],
-        files: [{ filename: "screenshot.png", value: ss.response.body }]
+        files: [{ filename: "screenshot.png", value: Buffer.from(ss.response.body.image, 'base64') }]
       })
     } catch (e) {
       console.log(e)
diff --git a/labscore/api/obelisk/endpoints.js b/labscore/api/obelisk/endpoints.js
index aebe8b8..fe570c2 100644
--- a/labscore/api/obelisk/endpoints.js
+++ b/labscore/api/obelisk/endpoints.js
@@ -11,7 +11,9 @@ const ObeliskApi = Object.freeze({
   LLM_MODELS_GENERATE: "/llm/v1/generate",
 
   GENERATIVEIMAGES_MODELS_IMAGEN: "/generativeimages/v1/models/imagen",
-  
+
+  WEBUTILS_WEBPAGE_SCREENSHOT: "/webutils/v1/webpage:screenshot",
+
   SPARK_WEB_SUMMARIZE: "/spark-pa/v1/web:summarize",
 
   WOLFRAM_QUERY_COMPUTE: "/wolfram/v1/query:compute",
diff --git a/labscore/api/obelisk/index.js b/labscore/api/obelisk/index.js
index fc53878..5ce2958 100644
--- a/labscore/api/obelisk/index.js
+++ b/labscore/api/obelisk/index.js
@@ -62,6 +62,13 @@ module.exports.GenerativeImagesModelsImagen = async function(context, prompt){
   })
 }
 
+module.exports.WebUtilsWebPageScreenshot = async function(context, url, allow_adult){
+  return await request(ObeliskApi.WEBUTILS_WEBPAGE_SCREENSHOT, "POST", {}, {
+    url,
+    allow_adult
+  })
+}
+
 module.exports.SparkWebSummarize = async function(context, url){
   return await request(ObeliskApi.SPARK_WEB_SUMMARIZE, "POST", {}, {
     url