feat(core/utils/puppeteer-utils): Cookie utils (#11471)

Signed-off-by: Rongrong <i@rong.moe>

Signed-off-by: Rongrong <i@rong.moe>

lib/utils/puppeteer-utils.js

@@ -0,0 +1,63 @@
+/**
+ * Get Cookie-header-style cookie string from a puppeteer-style cookie array
+ *
+ * @param {import('puppeteer').Protocol.Network.CookieParam[]} cookies Puppeteer-style cookie array
+ * @param {RegExp | string} domainFilter Filter cookies by domain or RegExp
+ * @return {string} Cookie-header-style cookie string (e.g. "foobar; foo=bar; baz=qux")
+ */
+const parseCookieArray = (cookies, domainFilter = null) => {
+    if (typeof domainFilter === 'string') {
+        const dotDomain = '.' + domainFilter;
+        cookies = cookies.filter(({ domain }) => domain === domainFilter || domain.endsWith(dotDomain));
+    } else if (domainFilter && domainFilter.test !== undefined) {
+        cookies = cookies.filter(({ domain }) => domainFilter.test(domain));
+    }
+    // {name: '', value: 'foobar'} => 'foobar' // https://stackoverflow.com/questions/42531198/cookie-without-a-name
+    // {name: 'foo', value: 'bar'} => 'foo=bar'
+    return cookies.map(({ name, value }) => (name ? `${name}=${value}` : value)).join('; ');
+};
+
+/**
+ * Construct a puppeteer-style cookie array from a Cookie-header-style cookie string
+ *
+ * @param {string} cookieStr Cookie-header-style cookie string (e.g. "foobar; foo=bar; baz=qux")
+ * @param {string} domain Domain to set for each cookie
+ * @return {import('puppeteer').Protocol.Network.CookieParam[]} Puppeteer-style cookie array
+ */
+const constructCookieArray = (cookieStr, domain) =>
+    cookieStr.split('; ').map((item) => {
+        const [name, value] = item.split('=');
+        return value === undefined ? { name: '', value: name, domain } : { name, value, domain };
+    });
+
+/**
+ * Set cookies for a page
+ *
+ * @param {import('puppeteer').Page} page Puppeteer Page object
+ * @param {string} cookieStr Cookie-header-style cookie string (e.g. "foobar; foo=bar; baz=qux")
+ * @param {string} domain Domain to set for each cookie
+ * @return {Promise<void>}
+ */
+const setCookies = async (page, cookieStr, domain) => {
+    const cookies = constructCookieArray(cookieStr, domain);
+    await page.setCookie(...cookies);
+};
+
+/**
+ * Get Cookie-header-style cookie string from a page
+ *
+ * @param {import('puppeteer').Page} page Puppeteer Page object
+ * @param {RegExp | string} domainFilter Filter cookies by domain or RegExp
+ * @return {Promise<string>} Cookie-header-style cookie string
+ */
+const getCookies = async (page, domainFilter = null) => {
+    const cookies = await page.cookies();
+    return parseCookieArray(cookies, domainFilter);
+};
+
+module.exports = {
+    parseCookieArray,
+    constructCookieArray,
+    setCookies,
+    getCookies,
+};

lib/v2/nuaa/utils/pypasswaf.js

@@ -1,3 +1,5 @@
+const puppeteerUtils = require('@/utils/puppeteer-utils');
+
 /**
  * async function 获取cookie
  * @desc 返回一个可用的cookie，使用 `got` 发起请求的时候，传入到`options.headers.cookie`即可
@@ -14,8 +16,7 @@ module.exports = async function getCookie(host) {
         waitUntil: 'networkidle0',
     });
 
-    let cookie = await page.cookies();
+    const cookie = await puppeteerUtils.getCookies(page);
-    cookie = cookie.map(({ name, value }) => `${name}=${value}`).join('; ');
     await browser.close();
     return cookie;
 };

test/utils/puppeteer-utils.js

@@ -0,0 +1,102 @@
+let puppeteer;
+const { parseCookieArray, constructCookieArray, setCookies, getCookies } = require('../../lib/utils/puppeteer-utils');
+
+let browser = null;
+
+afterEach(() => {
+    if (browser) {
+        browser.close();
+        browser = null;
+    }
+
+    jest.resetModules();
+});
+
+describe('puppeteer-utils', () => {
+    const cookieArrayExampleCom = [
+        { name: 'foobar', value: '', domain: 'example.com' },
+        { name: 'foo', value: 'bar', domain: 'example.com' },
+        { name: 'baz', value: 'qux', domain: 'example.com' },
+    ];
+    const cookieArraySubExampleCom = [
+        { name: 'barfoo', value: '', domain: 'sub.example.com' },
+        { name: 'bar', value: 'foo', domain: 'sub.example.com' },
+        { name: 'qux', value: 'baz', domain: 'sub.example.com' },
+    ];
+    const cookieArrayRsshubTest = [
+        { name: '', value: 'rsshub', domain: 'rsshub.test' },
+        { name: 'rsshub', value: '', domain: 'rsshub.test' },
+        { name: 'test', value: 'rsshub', domain: 'rsshub.test' },
+    ];
+    const cookieArrayAll = cookieArrayExampleCom.concat(cookieArraySubExampleCom).concat(cookieArrayRsshubTest);
+
+    const cookieStrExampleCom = 'foobar=; foo=bar; baz=qux';
+    const cookieStrSubExampleCom = 'barfoo=; bar=foo; qux=baz';
+    const cookieStrRsshubTest = 'rsshub; rsshub=; test=rsshub';
+    const cookieStrAll = [cookieStrExampleCom, cookieStrSubExampleCom, cookieStrRsshubTest].join('; ');
+
+    it('parseCookieArray', () => {
+        [
+            [cookieArrayExampleCom, cookieStrExampleCom],
+            [cookieArraySubExampleCom, cookieStrSubExampleCom],
+            [cookieArrayRsshubTest, cookieStrRsshubTest],
+            [cookieArrayAll, cookieStrAll],
+        ].forEach(([cookieArray, cookieStr]) => {
+            expect(parseCookieArray(cookieArray)).toEqual(cookieStr);
+        });
+        expect(parseCookieArray(cookieArrayAll, 'example.com')).toEqual(`${cookieStrExampleCom}; ${cookieStrSubExampleCom}`);
+        expect(parseCookieArray(cookieArrayAll, 'sub.example.com')).toEqual(cookieStrSubExampleCom);
+        expect(parseCookieArray(cookieArrayExampleCom, 'sub.example.com')).toEqual('');
+        expect(parseCookieArray(cookieArrayRsshubTest, 'example.com')).toEqual('');
+        expect(parseCookieArray(cookieArraySubExampleCom, 'example.com')).toEqual(cookieStrSubExampleCom);
+        expect(parseCookieArray(cookieArrayAll, 'rsshub.test')).toEqual(cookieStrRsshubTest);
+        expect(parseCookieArray(cookieArrayAll, /^example\.com$/)).toEqual(cookieStrExampleCom);
+        expect(parseCookieArray(cookieArrayAll, /^sub\.example\.com|rsshub\.test$/)).toEqual(`${cookieStrSubExampleCom}; ${cookieStrRsshubTest}`);
+        expect(parseCookieArray(cookieArrayAll, /^.*$/)).toEqual(cookieStrAll);
+    });
+
+    it('constructCookieArray', () => {
+        [
+            [cookieArrayExampleCom, cookieStrExampleCom],
+            [cookieArraySubExampleCom, cookieStrSubExampleCom],
+            [cookieArrayRsshubTest, cookieStrRsshubTest],
+        ].forEach(([cookieArray, cookieStr]) => {
+            expect(constructCookieArray(cookieStr, cookieArray[0].domain)).toEqual(cookieArray);
+        });
+    });
+
+    it('getCookies httpbin', async () => {
+        puppeteer = require('../../lib/utils/puppeteer');
+        browser = await puppeteer();
+        const page = await browser.newPage();
+        await page.goto('https://httpbin.org/cookies/set?foo=bar&baz=qux', {
+            waitUntil: 'domcontentloaded',
+        });
+        expect((await getCookies(page, 'httpbin.org')).split('; ').sort()).toEqual(['foo=bar', 'baz=qux'].sort());
+    }, 10000);
+
+    it('setCookies httpbin', async () => {
+        puppeteer = require('../../lib/utils/puppeteer');
+        browser = await puppeteer();
+        const page = await browser.newPage();
+        // httpbin.org cannot recognize cookies with empty name properly, so we cannot use cookieStrAll here
+        await setCookies(page, cookieStrExampleCom, 'httpbin.org');
+        await page.goto('https://httpbin.org/cookies', {
+            waitUntil: 'domcontentloaded',
+        });
+        const data = await page.evaluate(() => JSON.parse(document.body.innerText));
+        expect(data.cookies).toEqual(Object.fromEntries(cookieArrayExampleCom.map(({ name, value }) => [name, value])));
+    }, 10000);
+
+    it('setCookies & getCookies example.org', async () => {
+        puppeteer = require('../../lib/utils/puppeteer');
+        browser = await puppeteer();
+        const page = await browser.newPage();
+        // we can use cookieStrAll here!
+        await setCookies(page, cookieStrAll, 'example.org');
+        await page.goto('https://example.org', {
+            waitUntil: 'domcontentloaded',
+        });
+        expect((await getCookies(page, 'example.org')).split('; ').sort()).toEqual(cookieStrAll.split('; ').sort());
+    }, 10000);
+});