feat: add screenshot api

actions/update-user-name.ts

@@ -1,9 +1,10 @@
 "use server";
 
+import { revalidatePath } from "next/cache";
 import { auth } from "@/auth";
+
 import { prisma } from "@/lib/db";
 import { userNameSchema } from "@/lib/validations/user";
-import { revalidatePath } from "next/cache";
 
 export type FormData = {
   name: string;
@@ -11,7 +12,7 @@ export type FormData = {
 
 export async function updateUserName(userId: string, data: FormData) {
   try {
-    const session = await auth()
+    const session = await auth();
 
     if (!session?.user || session?.user.id !== userId) {
       throw new Error("Unauthorized");
@@ -27,12 +28,12 @@ export async function updateUserName(userId: string, data: FormData) {
       data: {
         name: name,
       },
-    })
+    });
 
-    revalidatePath('/dashboard/settings');
+    revalidatePath("/dashboard/settings");
     return { status: "success" };
   } catch (error) {
     // console.log(error)
-    return { status: "error" }
+    return { status: "error" };
   }
-}
+}

app/(protected)/dashboard/scrape/meta-scraping.tsx

@@ -1,11 +1,11 @@
 "use client";
 
 import { useState } from "react";
+import { User } from "@prisma/client";
 import JsonView from "@uiw/react-json-view";
 import { githubLightTheme } from "@uiw/react-json-view/githubLight";
 import { vscodeTheme } from "@uiw/react-json-view/vscode";
 import { useTheme } from "next-themes";
-import puppeteer from "puppeteer";
 import { toast } from "sonner";
 
 import { Button } from "@/components/ui/button";
@@ -37,7 +37,11 @@ export interface MetaScrapingProps {
   timestamp: string;
 }
 
-export default function MetaScraping() {
+export default function MetaScraping({
+  user,
+}: {
+  user: { id: string; apiKey: string };
+}) {
   const { theme } = useTheme();
   const [currentLink, setCurrentLink] = useState("wr.do");
   const [protocol, setProtocol] = useState("https://");
@@ -56,16 +60,14 @@ export default function MetaScraping() {
   const [isShoting, setIsShoting] = useState(false);
   const [currentScreenshotLink, setCurrentScreenshotLink] = useState("wr.do");
   const [screenshotInfo, setScreenshotInfo] = useState({
-    data: "",
     url: "",
-    timestamp: "",
   });
 
   const handleScrapingMeta = async () => {
     if (currentLink) {
       setIsScraping(true);
       const res = await fetch(
-        `/api/scraping/meta?url=${protocol}${currentLink}`,
+        `/api/scraping/meta?url=${protocol}${currentLink}&key=${user.apiKey}`,
       );
       if (!res.ok || res.status !== 200) {
         toast.error(res.statusText || "Error");
@@ -79,21 +81,21 @@ export default function MetaScraping() {
   };
 
   const handleScrapingScreenshot = async () => {
-    // if (currentScreenshotLink) {
-    //   setIsShoting(true);
-    //   const res = await fetch(
-    //     `/api/scraping/screenshot?url=${protocol}${currentScreenshotLink}`,
-    //   );
-    //   if (!res.ok || res.status !== 200) {
-    //     toast.error(res.statusText);
-    //   } else {
-    //     const data = await res.json();
-    //     console.log(data);
-    //     setScreenshotInfo(data);
-    //     toast.success("Success!");
-    //   }
-    //   setIsShoting(false);
-    // }
+    if (currentScreenshotLink) {
+      setIsShoting(true);
+      const res = await fetch(
+        `/api/scraping/screenshot?url=${protocol}${currentScreenshotLink}&key=${user.apiKey}`,
+      );
+      if (!res.ok || res.status !== 200) {
+        toast.error(res.statusText);
+      } else {
+        const blob = await res.blob();
+        const imageUrl = URL.createObjectURL(blob);
+        setScreenshotInfo({ url: imageUrl });
+        toast.success("Success!");
+      }
+      setIsShoting(false);
+    }
   };
 
   return (
@@ -152,7 +154,7 @@ export default function MetaScraping() {
           </div>
         </CardContent>
       </Card>
-      {/* <Card>
+      <Card>
         <CardHeader>
           <CardTitle>Screenshot</CardTitle>
           <CardDescription>
@@ -208,9 +210,9 @@ export default function MetaScraping() {
               displayDataTypes={false}
               // shortenTextAfterLength={50}
             />
-            {screenshotInfo.data && (
+            {screenshotInfo.url && (
               <BlurImage
-                src={screenshotInfo.data}
+                src={screenshotInfo.url}
                 alt="ligth preview landing"
                 className="my-4 flex rounded-md border object-contain object-center shadow-md dark:hidden"
                 width={1500}
@@ -221,7 +223,7 @@ export default function MetaScraping() {
             )}
           </div>
         </CardContent>
-      </Card> */}
+      </Card>
     </>
   );
 }

app/(protected)/dashboard/scrape/page.tsx

@@ -24,7 +24,7 @@ export default async function DashboardPage() {
         link="/docs/scraping-api"
         linkText="Scraping API."
       />
-      <MetaScraping />
+      <MetaScraping user={{ id: user.id, apiKey: user.apiKey }} />
     </>
   );
 }

app/(protected)/dashboard/settings/page.tsx

@@ -4,6 +4,7 @@ import { getCurrentUser } from "@/lib/session";
 import { constructMetadata } from "@/lib/utils";
 import { DeleteAccountSection } from "@/components/dashboard/delete-account";
 import { DashboardHeader } from "@/components/dashboard/header";
+import { UserApiKeyForm } from "@/components/forms/user-api-key-form";
 import { UserNameForm } from "@/components/forms/user-name-form";
 import { UserRoleForm } from "@/components/forms/user-role-form";
 
@@ -28,6 +29,13 @@ export default async function SettingsPage() {
         {user.role === "ADMIN" && (
           <UserRoleForm user={{ id: user.id, role: user.role }} />
         )}
+        <UserApiKeyForm
+          user={{
+            id: user.id,
+            name: user.name || "",
+            apiKey: user.apiKey || "",
+          }}
+        />
         <DeleteAccountSection />
       </div>
     </>

app/api/keys/route.ts

@@ -0,0 +1,25 @@
+import { env } from "@/env.mjs";
+import { generateApiKey } from "@/lib/dto/api-key";
+import { checkUserStatus } from "@/lib/dto/user";
+import { getCurrentUser } from "@/lib/session";
+
+export async function POST(req: Request) {
+  try {
+    const user = checkUserStatus(await getCurrentUser());
+    if (user instanceof Response) return user;
+
+    const res = await generateApiKey(user.id);
+    if (res) {
+      return Response.json(res.apiKey);
+    }
+    return Response.json(res, {
+      status: 501,
+      statusText: "Server error",
+    });
+  } catch (error) {
+    return Response.json("An error occurred", {
+      status: error.status || 500,
+      statusText: error.statusText || "Server error",
+    });
+  }
+}

app/api/record/add/route.ts

@@ -28,9 +28,10 @@ export async function POST(req: Request) {
       });
     }
 
-    // check quota
+    // Check quota: 若是管理员则不检查，否则检查
     const user_records_count = await getUserRecordCount(user.id);
     if (
+      user.role !== "ADMIN" &&
       Number(NEXT_PUBLIC_FREE_RECORD_QUOTA) > 0 &&
       user_records_count >= Number(NEXT_PUBLIC_FREE_RECORD_QUOTA)
     ) {

app/api/scraping/meta/route.ts

@@ -1,5 +1,6 @@
 import cheerio from "cheerio";
 
+import { checkApiKey } from "@/lib/dto/api-key";
 import { checkUserStatus } from "@/lib/dto/user";
 import { getCurrentUser } from "@/lib/session";
 import { isLink, removeUrlSuffix } from "@/lib/utils";
@@ -8,9 +9,6 @@ export const revalidate = 600;
 
 export async function GET(req: Request) {
   try {
-    const user = checkUserStatus(await getCurrentUser());
-    if (user instanceof Response) return user;
-
     const url = new URL(req.url);
     const link = url.searchParams.get("url");
     if (!link || !isLink(link)) {
@@ -20,6 +18,24 @@ export async function GET(req: Request) {
       });
     }
 
+    // Get the API key from the request
+    const custom_apiKey = url.searchParams.get("key");
+    if (!custom_apiKey) {
+      return Response.json("API key is required", {
+        status: 400,
+        statusText: "API key is required",
+      });
+    }
+
+    // Check if the API key is valid
+    const user_apiKey = await checkApiKey(custom_apiKey);
+    if (!user_apiKey?.id) {
+      return Response.json("Invalid API key", {
+        status: 403,
+        statusText: "Invalid API key",
+      });
+    }
+
     const res = await fetch(link);
     if (!res.ok) {
       return Response.json("Failed to fetch url", {
@@ -69,7 +85,7 @@ export async function GET(req: Request) {
       timestamp: Date.now(),
     });
   } catch (error) {
-    console.log(error);
+    // console.log(error);
     return Response.json("An error occurred", {
       status: error.status || 500,
       statusText: error.statusText || "Server error",

app/api/scraping/screenshot/route.ts

@@ -1,5 +1,5 @@
-// import puppeteer from "puppeteer";
-
+import { env } from "@/env.mjs";
+import { checkApiKey } from "@/lib/dto/api-key";
 import { checkUserStatus } from "@/lib/dto/user";
 import { getCurrentUser } from "@/lib/session";
 import { isLink } from "@/lib/utils";
@@ -7,47 +7,72 @@ import { isLink } from "@/lib/utils";
 export const revalidate = 60;
 
 // export const runtime = "edge";
-const puppeteer = require("puppeteer");
 
 export async function GET(req: Request) {
-  const user = checkUserStatus(await getCurrentUser());
-  if (user instanceof Response) return user;
-
-  const url = new URL(req.url);
-  const link = url.searchParams.get("url");
-  const full = url.searchParams.get("full") || "false";
-
-  if (!link || !isLink(link)) {
-    return Response.json("Invalid url", {
-      status: 400,
-      statusText: "Invalid url",
-    });
-  }
-
-  let browser;
-
   try {
-    browser = await puppeteer.launch();
-    const page = await browser.newPage();
-    await page.setViewport({ width: 1600, height: 1200 });
-    await page.goto(link);
-    const screenshot = await page.screenshot({
-      encoding: "base64",
-      fullPage: full === "true",
-    });
-    return Response.json({
-      data: `data:image/png;base64,${screenshot}`,
-      url: link,
-      timestamp: Date.now(),
+    const url = new URL(req.url);
+    const link = url.searchParams.get("url");
+    const full = url.searchParams.get("full") || "false";
+    const width = url.searchParams.get("width") || "1600";
+    const height = url.searchParams.get("height") || "1200";
+    const viewportWidth = url.searchParams.get("viewportWidth") || "1080";
+    const viewportHeight = url.searchParams.get("viewportHeight") || "1080";
+    const forceReload = url.searchParams.get("forceReload") || "false";
+    const isMobile = url.searchParams.get("isMobile") || "false";
+    const isDarkMode = url.searchParams.get("isDarkMode") || "false";
+    const deviceScaleFactor = url.searchParams.get("deviceScaleFactor") || "1";
+
+    // Check if the url is valid
+    if (!link || !isLink(link)) {
+      return Response.json("Invalid url", {
+        status: 400,
+        statusText: "Invalid url",
+      });
+    }
+
+    // Get the API key from the request
+    const custom_apiKey = url.searchParams.get("key");
+    if (!custom_apiKey) {
+      return Response.json("API key is required", {
+        status: 400,
+        statusText: "API key is required",
+      });
+    }
+
+    // Check if the API key is valid
+    const user_apiKey = await checkApiKey(custom_apiKey);
+    if (!user_apiKey?.id) {
+      return Response.json("Invalid API key", {
+        status: 403,
+        statusText: "Invalid API key",
+      });
+    }
+
+    const { SCREENSHOTONE_BASE_URL } = env;
+    const scrape_url = `${SCREENSHOTONE_BASE_URL}?url=${link}&isFullPage=${full}&width=${width}&height=${height}&viewportWidth=${viewportWidth}&viewportHeight=${viewportHeight}&forceReload=${forceReload}&isMobile=${isMobile}&isDarkMode=${isDarkMode}&deviceScaleFactor=${deviceScaleFactor}`;
+    console.log("[Scrape Url]", scrape_url);
+
+    const res = await fetch(scrape_url);
+    if (!res.ok) {
+      return Response.json("Failed to get screenshot", {
+        status: 406,
+        statusText: "Failed to get screenshot",
+      });
+    }
+
+    const imageBuffer = await res.arrayBuffer();
+    return new Response(imageBuffer, {
+      status: 200,
+      headers: {
+        "Content-Type": "image/png",
+        "Cache-Control": "public, max-age=86400, immutable",
+      },
     });
   } catch (error) {
-    return Response.json(error, {
-      status: 500,
-      statusText: error,
+    console.log(error);
+    return Response.json("Server error", {
+      status: error.status || 500,
+      statusText: "Server error",
     });
-  } finally {
-    if (browser) {
-      await browser.close();
-    }
   }
 }

auth.ts

@@ -13,6 +13,7 @@ declare module "next-auth" {
       role: UserRole;
       team: string;
       active: number;
+      apiKey: string;
     } & DefaultSession["user"];
   }
 }
@@ -46,6 +47,7 @@ export const {
         session.user.image = token.picture;
         session.user.active = token.active as number;
         session.user.team = token.team as string;
+        session.user.apiKey = token.apiKey as string;
       }
 
       return session;
@@ -64,6 +66,7 @@ export const {
       token.role = dbUser.role;
       token.active = dbUser.active;
       token.team = dbUser.team;
+      token.apiKey = dbUser.apiKey;
 
       return token;
     },

components/forms/user-api-key-form.tsx

@@ -0,0 +1,105 @@
+"use client";
+
+import { useState, useTransition } from "react";
+import { zodResolver } from "@hookform/resolvers/zod";
+import { User } from "@prisma/client";
+import { useForm } from "react-hook-form";
+import { toast } from "sonner";
+
+import { cn } from "@/lib/utils";
+import { userApiKeySchema } from "@/lib/validations/user";
+import { Button } from "@/components/ui/button";
+import { Input } from "@/components/ui/input";
+import { Label } from "@/components/ui/label";
+import { SectionColumns } from "@/components/dashboard/section-columns";
+import { Icons } from "@/components/shared/icons";
+
+import { CopyButton } from "../shared/copy-button";
+
+interface UserNameFormProps {
+  user: Pick<User, "id" | "name" | "apiKey">;
+}
+
+type FormData = {
+  apiKey: string;
+};
+
+export function UserApiKeyForm({ user }: UserNameFormProps) {
+  const [isPending, startTransition] = useTransition();
+  const [apiKey, setApiKey] = useState(user?.apiKey || "");
+
+  const {
+    handleSubmit,
+    formState: { errors },
+  } = useForm<FormData>({
+    defaultValues: {
+      apiKey: user?.apiKey || "",
+    },
+  });
+
+  const onSubmit = handleSubmit(() => {
+    startTransition(async () => {
+      const response = await fetch(`/api/keys`, {
+        method: "POST",
+      });
+      if (!response.ok || response.status !== 200) {
+        toast.error("Created Failed!", {
+          description: response.statusText,
+        });
+      } else {
+        const res = await response.json();
+        setApiKey(res);
+        toast.success(`Generated successfully!`);
+      }
+    });
+  });
+
+  return (
+    <form onSubmit={onSubmit}>
+      <SectionColumns
+        title="API Key"
+        description="Generate a new API key to access the scraper apis."
+      >
+        <div className="flex w-full items-center gap-2">
+          <Label className="sr-only" htmlFor="name">
+            API Key
+          </Label>
+          <div className="flex w-full items-center">
+            <input
+              value={apiKey || "Click to generate your API key"}
+              disabled
+              className="flex h-9 flex-1 shrink-0 items-center truncate rounded-l-md border border-r-0 border-input bg-transparent px-3 py-2 text-sm"
+            />
+            <CopyButton
+              value={apiKey}
+              className={cn(
+                "size-[36px]",
+                "duration-250 rounded-l-none border transition-all group-hover:opacity-100",
+              )}
+            />
+          </div>
+
+          <Button
+            type="submit"
+            variant={"blue"}
+            disabled={isPending}
+            className="shrink-0 px-4"
+          >
+            {isPending ? (
+              <Icons.spinner className="size-4 animate-spin" />
+            ) : (
+              <p>Generate</p>
+            )}
+          </Button>
+        </div>
+        <div className="flex flex-col justify-between p-1">
+          {errors?.apiKey && (
+            <p className="pb-0.5 text-[13px] text-red-600">
+              {errors.apiKey.message}
+            </p>
+          )}
+        </div>
+      </SectionColumns>
+    </form>
+  );
+}

components/forms/user-auth-form.tsx

@@ -59,7 +59,7 @@ export function UserAuthForm({ className, type, ...props }: UserAuthFormProps) {
 
   return (
     <div className={cn("grid gap-3", className)} {...props}>
-      {/* <form onSubmit={handleSubmit(onSubmit)}>
+      <form onSubmit={handleSubmit(onSubmit)}>
         <div className="grid gap-2">
           <div className="grid gap-1">
             <Label className="sr-only" htmlFor="email">
@@ -107,7 +107,7 @@ export function UserAuthForm({ className, type, ...props }: UserAuthFormProps) {
             Or continue with
           </span>
         </div>
-      </div> */}
+      </div>
 
       <button
         type="button"

components/shared/copy-button.tsx

@@ -27,6 +27,7 @@ export function CopyButton({ value, className, ...props }: CopyButtonProps) {
 
   return (
     <Button
+      type="button"
       size="sm"
       variant="ghost"
       className={cn(

env.mjs

@@ -16,6 +16,7 @@ export const env = createEnv({
     CLOUDFLARE_ZONE_ID: z.string().min(1),
     CLOUDFLARE_API_KEY: z.string().min(1),
     CLOUDFLARE_EMAIL: z.string().min(1),
+    SCREENSHOTONE_BASE_URL: z.string().min(1),
   },
   client: {
     NEXT_PUBLIC_APP_URL: z.string().min(1),
@@ -39,5 +40,6 @@ export const env = createEnv({
     CLOUDFLARE_ZONE_ID: process.env.CLOUDFLARE_ZONE_ID,
     CLOUDFLARE_API_KEY: process.env.CLOUDFLARE_API_KEY,
     CLOUDFLARE_EMAIL: process.env.CLOUDFLARE_EMAIL,
+    SCREENSHOTONE_BASE_URL: process.env.SCREENSHOTONE_BASE_URL,
   },
 });

lib/dto/api-key.ts

@@ -0,0 +1,26 @@
+import crypto from "crypto";
+
+import { prisma } from "@/lib/db";
+
+export const getApiKeyByUserId = async (userId: string) => {
+  return prisma.user.findUnique({
+    where: { id: userId },
+    select: { apiKey: true },
+  });
+};
+
+export const checkApiKey = async (apiKey: string) => {
+  return prisma.user.findFirst({
+    where: { apiKey },
+    select: { id: true },
+  });
+};
+
+export const generateApiKey = async (userId: string) => {
+  const apiKey = crypto.randomUUID();
+  return prisma.user.update({
+    where: { id: userId },
+    data: { apiKey },
+    select: { apiKey: true },
+  });
+};

lib/validations/user.ts

@@ -8,3 +8,5 @@ export const userNameSchema = z.object({
 export const userRoleSchema = z.object({
   role: z.nativeEnum(UserRole),
 });
+
+export const userApiKeySchema = z.object({});

package.json

@@ -79,7 +79,6 @@
     "next-view-transitions": "^0.3.0",
     "nodemailer": "^6.9.14",
     "prop-types": "^15.8.1",
-    "puppeteer": "^23.0.1",
     "react": "18.3.1",
     "react-countup": "^6.5.3",
     "react-day-picker": "^8.10.1",

prisma/migrations/20240705091917_init/migration.sql

@@ -175,4 +175,6 @@ ALTER TABLE "url_metas" ADD COLUMN "lang" TEXT;
 ALTER TABLE "url_metas" ADD COLUMN "device" TEXT;
 ALTER TABLE "url_metas" ADD COLUMN "browser" TEXT;
 
-ALTER TABLE "user_urls" ADD COLUMN "expiration" TEXT NOT NULL DEFAULT '-1';
+ALTER TABLE "user_urls" ADD COLUMN "expiration" TEXT NOT NULL DEFAULT '-1';
+
+ALTER TABLE "users" ADD COLUMN "apiKey" TEXT;

prisma/schema.prisma

@@ -57,6 +57,7 @@ model User {
   image         String?
   active        Int       @default(1) // 0 封禁，1 正常
   team          String?
+  apiKey        String?
   createdAt     DateTime  @default(now()) @map(name: "created_at")
   updatedAt     DateTime  @default(now()) @map(name: "updated_at")
   role          UserRole  @default(USER)