#1958 added support for generic url downloading (current whitelist: discordapp.com, catbox.moe)

This commit is contained in:
Adam 2024-04-29 23:09:21 -05:00 committed by Adam
parent d9d76ba16d
commit 300b68177b
1 changed files with 74 additions and 4 deletions

View File

@ -426,6 +426,32 @@ function parseAICC(url) {
return null; 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 * @param {String} url
* @returns {String | null } UUID of the character * @returns {String | null } UUID of the character
@ -440,6 +466,41 @@ function getUuidFromUrl(url) {
return uuid; return uuid;
} }
/**
* Filter to get the domain host of a url instead of a blanket string
* search.
* @param {String} url URL to strip
* @result {String} Domain name
*/
function getHostFromUrl(url) {
//dirty way to do it, but if someone wants to bash their head on the regex
//go ahead
return url.replace('http://','').replace('https://','').split('/')[0];
}
/**
* List of generic domains that are trusted to hold character cards.
* Some additional providers may be added later on.
*/
const WHITELIST_GENERIC_URL_DOWNLOAD_SOURCES = [
'discordapp.com',
'catbox.moe',
]; //TODO: move this to config.yaml?
/**
* Checks if url is part of generic download source whitelist.
* @param {String} url
* @returns {bool} if the url is on the whitelist.
*/
function isUrlWhitelisted(url) {
for(const k of WHITELIST_GENERIC_URL_DOWNLOAD_SOURCES) {
if(url.includes(k)) {
return true;
}
}
return false;
}
const router = express.Router(); const router = express.Router();
router.post('/importURL', jsonParser, async (request, response) => { router.post('/importURL', jsonParser, async (request, response) => {
@ -449,12 +510,15 @@ router.post('/importURL', jsonParser, async (request, response) => {
try { try {
const url = request.body.url; const url = request.body.url;
const host = getHostFromUrl(url);
let result; let result;
let type; let type;
const isJannnyContent = url.includes('janitorai'); const isChub = host.includes('chub.ai');
const isPygmalionContent = url.includes('pygmalion.chat'); const isJannnyContent = host.includes('janitorai');
const isAICharacterCardsContent = url.includes('aicharactercards.com'); const isPygmalionContent = host.includes('pygmalion.chat');
const isAICharacterCardsContent = host.includes('aicharactercards.com');
const isGeneric = isUrlWhitelisted(host);
if (isPygmalionContent) { if (isPygmalionContent) {
const uuid = getUuidFromUrl(url); const uuid = getUuidFromUrl(url);
@ -479,7 +543,7 @@ router.post('/importURL', jsonParser, async (request, response) => {
} }
type = 'character'; type = 'character';
result = await downloadAICCCharacter(AICCParsed); result = await downloadAICCCharacter(AICCParsed);
} else { } else if (isChub) {
const chubParsed = parseChubUrl(url); const chubParsed = parseChubUrl(url);
type = chubParsed?.type; type = chubParsed?.type;
@ -494,6 +558,12 @@ router.post('/importURL', jsonParser, async (request, response) => {
else { else {
return response.sendStatus(404); 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.fileType) response.set('Content-Type', result.fileType); if (result.fileType) response.set('Content-Type', result.fileType);