Merge pull request #2171 from 24adamcho/generic-card-download

Character card import from generic sources (specifically Discord, Catbox.moe)

default/config.yaml

@@ -48,6 +48,12 @@ allowKeysExposure: false
 skipContentCheck: false
 # Disable automatic chats backup
 disableChatBackup: false
+# Allowed hosts for card downloads
+whitelistImportDomains:
+  - localhost
+  - cdn.discordapp.com
+  - files.catbox.moe
+  - raw.githubusercontent.com
 # API request overrides (for KoboldAI and Text Completion APIs)
 ## Note: host includes the port number if it's not the default (80 or 443)
 ## Format is an array of objects:

public/script.js

@@ -10165,7 +10165,7 @@ jQuery(async function () {
             <li>JanitorAI Character (Direct Link or UUID)<br>Example: <tt>ddd1498a-a370-4136-b138-a8cd9461fdfe_character-aqua-the-useless-goddess</tt></li>
             <li>Pygmalion.chat Character (Direct Link or UUID)<br>Example: <tt>a7ca95a1-0c88-4e23-91b3-149db1e78ab9</tt></li>
             <li>AICharacterCard.com Character (Direct Link or ID)<br>Example: <tt>AICC/aicharcards/the-game-master</tt></li>
-            <li>More coming soon...</li>
+            <li>Direct PNG Link (refer to <code>config.yaml</code> for allowed hosts)<br>Example: <tt>https://files.catbox.moe/notarealfile.png</tt></li>
         <ul>`;
         const input = await callPopup(html, 'input', '', { okButton: 'Import', rows: 4 });
 

src/endpoints/content-manager.js

@@ -10,6 +10,8 @@ const contentDirectory = path.join(process.cwd(), 'default/content');
 const contentIndexPath = path.join(contentDirectory, 'index.json');
 const characterCardParser = require('../character-card-parser.js');
 
+const WHITELIST_GENERIC_URL_DOWNLOAD_SOURCES = getConfigValue('whitelistImportDomains', []);
+
 /**
  * @typedef {Object} ContentItem
  * @property {string} filename
@@ -426,6 +428,32 @@ function parseAICC(url) {
     return null;
 }
 
+/**
+ * Download character card from generic url.
+ * @param {String} url
+ */
+async function downloadGenericPng(url) {
+    try {
+        const result = await fetch(url);
+
+        if (result.ok) {
+            const buffer = await result.buffer();
+            const fileName = sanitize(result.url.split('?')[0].split('/').reverse()[0]);
+            const contentType = result.headers.get('content-type') || 'image/png'; //yoink it from AICC function lol
+
+            return {
+                buffer: buffer,
+                fileName: fileName,
+                fileType: contentType,
+            };
+        }
+    } catch (error) {
+        console.error('Error downloading file: ', error);
+        throw error;
+    }
+    return null;
+}
+
 /**
 * @param {String} url
 * @returns {String | null } UUID of the character
@@ -440,6 +468,29 @@ function getUuidFromUrl(url) {
     return uuid;
 }
 
+/**
+ * Filter to get the domain host of a url instead of a blanket string search.
+ * @param {String} url URL to strip
+ * @returns {String} Domain name
+ */
+function getHostFromUrl(url) {
+    try {
+        const urlObj = new URL(url);
+        return urlObj.hostname;
+    } catch {
+        return '';
+    }
+}
+
+/**
+ * Checks if host is part of generic download source whitelist.
+ * @param {String} host Host to check
+ * @returns {boolean} If the host is on the whitelist.
+ */
+function isHostWhitelisted(host) {
+    return WHITELIST_GENERIC_URL_DOWNLOAD_SOURCES.includes(host);
+}
+
 const router = express.Router();
 
 router.post('/importURL', jsonParser, async (request, response) => {
@@ -449,12 +500,15 @@ router.post('/importURL', jsonParser, async (request, response) => {
 
     try {
         const url = request.body.url;
+        const host = getHostFromUrl(url);
         let result;
         let type;
 
-        const isJannnyContent = url.includes('janitorai');
-        const isPygmalionContent = url.includes('pygmalion.chat');
-        const isAICharacterCardsContent = url.includes('aicharactercards.com');
+        const isChub = host.includes('chub.ai');
+        const isJannnyContent = host.includes('janitorai');
+        const isPygmalionContent = host.includes('pygmalion.chat');
+        const isAICharacterCardsContent = host.includes('aicharactercards.com');
+        const isGeneric = isHostWhitelisted(host);
 
         if (isPygmalionContent) {
             const uuid = getUuidFromUrl(url);
@@ -479,7 +533,7 @@ router.post('/importURL', jsonParser, async (request, response) => {
             }
             type = 'character';
             result = await downloadAICCCharacter(AICCParsed);
-        } else {
+        } else if (isChub) {
             const chubParsed = parseChubUrl(url);
             type = chubParsed?.type;
 
@@ -494,10 +548,20 @@ router.post('/importURL', jsonParser, async (request, response) => {
             else {
                 return response.sendStatus(404);
             }
+        } else if (isGeneric) {
+            console.log('Downloading from generic url.');
+            type = 'character';
+            result = await downloadGenericPng(url);
+        } else {
+            return response.sendStatus(404);
+        }
+
+        if (!result) {
+            return response.sendStatus(404);
         }
 
         if (result.fileType) response.set('Content-Type', result.fileType);
-        response.set('Content-Disposition', `attachment; filename="${result.fileName}"`);
+        response.set('Content-Disposition', `attachment; filename="${encodeURI(result.fileName)}"`);
         response.set('X-Custom-Content-Type', type);
         return response.send(result.buffer);
     } catch (error) {