Actually made screenshot owner-only, made image API interaction more secure

api/index.js

@@ -39,7 +39,7 @@ if (isMainThread) {
     });
     worker.on("error", err => {
       console.error("worker error:", err);
-      socket.send(Buffer.concat([Buffer.from([0x2]), Buffer.from(err.toString())]), jobs[uuid].port, jobs[uuid].addr);
+      socket.send(Buffer.concat([Buffer.from([0x2]), Buffer.from(uuid), Buffer.from(err.toString())]), jobs[uuid].port, jobs[uuid].addr);
 
       workingWorkers--;
       if (queue.length > 0) {
@@ -143,10 +143,10 @@ if (isMainThread) {
           }, 500);
         }
       });
-      socket.send(Buffer.concat([Buffer.from([0x1]), Buffer.from(job.port.toString())]), job.port, job.addr);
+      socket.send(Buffer.concat([Buffer.from([0x1]), Buffer.from(job.uuid), Buffer.from(job.port.toString())]), job.port, job.addr);
       parentPort.postMessage(job.uuid); //Inform main thread about this worker freeing up
     } catch (e) {
-      socket.send(Buffer.concat([Buffer.from([0x2]), Buffer.from(e.toString())]), job.port, job.address);
+      socket.send(Buffer.concat([Buffer.from([0x2]), Buffer.from(job.uuid), Buffer.from(e.toString())]), job.port, job.address);
       parentPort.postMessage(job.uuid);
     }
   });

commands/screenshot.js

@@ -5,6 +5,7 @@ puppeteer.use(StealthPlugin());
 const fetch = require("node-fetch");
 
 exports.run = async (message, args) => {
+  if (message.author.id !== process.env.OWNER) return `${message.author.mention}, only the bot owner can run this command!`;
   message.channel.sendTyping();
   if (args.length === 0) return `${message.author.mention}, you need to provide a URL to screenshot!`;
   const getEndpoint = await fetch(`http://${process.env.CHROME}/json/version`);
@@ -28,5 +29,5 @@ exports.run = async (message, args) => {
 };
 
 exports.aliases = ["webshot", "ss", "shot", "page"];
-exports.category = 1;
+exports.category = 8;
 exports.help = "Screenshots a webpage";

utils/image.js

@@ -64,39 +64,45 @@ exports.run = (object, fromAPI = false) => {
       const currentServer = servers[Math.floor(Math.random() * servers.length)];
       const socket = dgram.createSocket("udp4");
       const data = Buffer.concat([Buffer.from([0x1]), Buffer.from(JSON.stringify(object))]);
-  
+
-      //let jobID;
+      let timeout = setTimeout(() => {
+        reject("Timed out");
+      }, 25000);
+      
+      let jobID;
       socket.on("message", (msg) => {
+        clearTimeout(timeout);
         const opcode = msg.readUint8(0);
-        const req = msg.slice(1, msg.length);
+        const req = msg.slice(37, msg.length);
+        const uuid = msg.slice(1, 36).toString();
         if (opcode === 0x0) {
-          //jobID = req;
+          jobID = uuid;
-          //console.log(`Our job UUID is: ${jobID}`);
+          timeout = setTimeout(() => {
+            reject("Timed out");
+          }, 300000);
         } else if (opcode === 0x1) {
-          //console.log(`Job ${jobID} is finished!`);
+          if (jobID === uuid) {
-          const client = net.createConnection(req.toString(), currentServer);
+            const client = net.createConnection(req.toString(), currentServer);
-          const array = [];
+            const array = [];
-          client.on("data", (rawData) => {
+            client.on("data", (rawData) => {
-            array.push(rawData);
+              array.push(rawData);
-            /*if (rawData.length < 32 * 1024) {
+            });
-              client.end();
+            client.once("end", () => {
-            }*/
+              const data = Buffer.concat(array);
-          });
+              const format = getFormat(data, "\n");
-          client.once("end", () => {
+              const payload = {
-            const data = Buffer.concat(array);
+                buffer: data.slice(format.dataStart + 1),
-            const format = getFormat(data, "\n");
+                type: format.buffer.toString().split("/")[1]
-            const payload = {
+              };
-              buffer: data.slice(format.dataStart + 1),
+              socket.close();
-              type: format.buffer.toString().split("/")[1]
+              resolve(payload);
-            };
+            });
-            socket.close();
+            client.on("error", (err) => {
-            resolve(payload);
+              throw err;
-          });
+            });
-          client.on("error", (err) => {
+          }
-            throw err;
-          });
         } else if (opcode === 0x2) {
-          reject(req);
+          if (jobID === uuid) reject(req);
         }
       });