From 9264d7afbe0a842daff12bf5176175f879b0a260 Mon Sep 17 00:00:00 2001 From: Rongrong Date: Tue, 20 Dec 2022 19:51:33 +0800 Subject: [PATCH] feat(core/utils/puppeteer-utils): Cookie utils (#11471) Signed-off-by: Rongrong Signed-off-by: Rongrong --- lib/utils/puppeteer-utils.js | 63 ++++++++++++++++++++ lib/v2/nuaa/utils/pypasswaf.js | 5 +- test/utils/puppeteer-utils.js | 102 +++++++++++++++++++++++++++++++++ 3 files changed, 168 insertions(+), 2 deletions(-) create mode 100644 lib/utils/puppeteer-utils.js create mode 100644 test/utils/puppeteer-utils.js diff --git a/lib/utils/puppeteer-utils.js b/lib/utils/puppeteer-utils.js new file mode 100644 index 000000000..44f00e8aa --- /dev/null +++ b/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} + */ +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} 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, +}; diff --git a/lib/v2/nuaa/utils/pypasswaf.js b/lib/v2/nuaa/utils/pypasswaf.js index cec40e60a..c17cc443f 100644 --- a/lib/v2/nuaa/utils/pypasswaf.js +++ b/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(); - cookie = cookie.map(({ name, value }) => `${name}=${value}`).join('; '); + const cookie = await puppeteerUtils.getCookies(page); await browser.close(); return cookie; }; diff --git a/test/utils/puppeteer-utils.js b/test/utils/puppeteer-utils.js new file mode 100644 index 000000000..90a18545f --- /dev/null +++ b/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); +});