import _ from 'lodash'; import { expect } from 'chai'; import { stub } from 'sinon'; import childProcess from 'child_process'; import * as firewall from '~/lib/firewall'; import * as iptables from '~/lib/iptables'; import { EventEmitter } from 'events'; import { Writable } from 'stream'; export enum RuleProperty { NotSet, } export type Testable = { [P in keyof T]?: T[P] | RuleProperty | undefined }; class FakeRuleAdaptor { private rules: iptables.Rule[]; constructor() { this.rules = []; } public getRuleAdaptor(): iptables.RuleAdaptor { return this.ruleAdaptor.bind(this); } private async ruleAdaptor(rules: iptables.Rule[]): Promise { const handleRule = async (rule: iptables.Rule) => { // remove any undefined values from the object... for (const key of Object.getOwnPropertyNames(rule)) { if ((rule as any)[key] === undefined) { delete (rule as any)[key]; } } this.rules.push(rule); return ''; }; if (Array.isArray(rules)) { for (const rule of rules) { await handleRule(rule); } } } private isSameRule( testable: Testable, rule: iptables.Rule, ): boolean { const props = Object.getOwnPropertyNames(testable); for (const prop of props) { if ( _.get(testable, prop) === RuleProperty.NotSet && _.get(rule, prop) === undefined ) { return true; } if ( _.get(rule, prop) === undefined || !_.isEqual(_.get(rule, prop), _.get(testable, prop)) ) { return false; } } return true; } public expectRule(testRule: Testable) { const matchingIndex = (() => { for (let i = 0; i < this.rules.length; i++) { if (this.isSameRule(testRule, this.rules[i])) { return i; } } return -1; })(); if (matchingIndex < 0) { console.log({ testRule, rules: this.rules }); } expect(matchingIndex).to.be.greaterThan(-1, `Rule has not been applied`); return matchingIndex; } public expectNoRule(testRule: Testable) { return expect( _.some(this.rules, (r) => this.isSameRule(testRule, r)), ).to.eq( false, `Rule has been applied: ${JSON.stringify(testRule)}\n\n${JSON.stringify( this.rules, null, 2, )}`, ); } public clearHistory() { this.rules = []; } } export const realRuleAdaptor = iptables.getDefaultRuleAdaptor(); const fakeRuleAdaptorManager = new FakeRuleAdaptor(); export const fakeRuleAdaptor = fakeRuleAdaptorManager.getRuleAdaptor(); // @ts-expect-error Assigning to a RO property iptables.getDefaultRuleAdaptor = () => { return fakeRuleAdaptor; }; export interface MockedState { hasAppliedRules: Promise; expectRule: (rule: Testable) => number; expectNoRule: (rule: Testable) => void; clearHistory: () => void; } export type MockedContext = (state: MockedState) => Promise; const applyFirewallRules = firewall.applyFirewallMode; export const whilstMocked = async ( context: MockedContext, ruleAdaptor: iptables.RuleAdaptor = fakeRuleAdaptor, ) => { const getOriginalDefaultRuleAdaptor = iptables.getDefaultRuleAdaptor; const spawnStub = stub(childProcess, 'spawn').callsFake(() => { const fakeProc = new EventEmitter(); (fakeProc as any).stdout = new EventEmitter(); const stdin = new Writable(); stdin._write = ( _chunk: Buffer, _encoding: string, callback: (err?: Error) => void, ) => { // console.log(chunk.toString('utf8')); callback(); fakeProc.emit('close', 1); }; (fakeProc as any).stdin = stdin; return fakeProc as any; }); // @ts-expect-error Assigning to a RO property iptables.getDefaultRuleAdaptor = () => { return ruleAdaptor; }; fakeRuleAdaptorManager.clearHistory(); const applied = new EventEmitter(); // @ts-expect-error Assigning to a RO property firewall.applyFirewallMode = async (mode: string) => { await applyFirewallRules(mode); applied.emit('applied'); }; await context({ expectRule: (rule) => fakeRuleAdaptorManager.expectRule(rule), expectNoRule: (rule) => fakeRuleAdaptorManager.expectNoRule(rule), clearHistory: () => fakeRuleAdaptorManager.clearHistory(), hasAppliedRules: new Promise((resolve) => { applied.once('applied', () => resolve()); }), }); // @ts-expect-error Assigning to a RO property firewall.applyFirewallMode = applyFirewallRules; spawnStub.restore(); // @ts-expect-error Assigning to a RO property iptables.getDefaultRuleAdaptor = getOriginalDefaultRuleAdaptor; };