Skip to content

Commit

Permalink
Merge pull request #9 from redcanaryco/github-timeouts
Browse files Browse the repository at this point in the history
GitHub timeouts
  • Loading branch information
rctgardner authored May 9, 2021
2 parents f599178 + 2a43c88 commit 5aa7a97
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 23 deletions.
32 changes: 21 additions & 11 deletions src/extension.ts
Original file line number Diff line number Diff line change
Expand Up @@ -51,12 +51,14 @@ export async function cacheData(storageDir: string): Promise<AttackMap|undefined
log('Checking extension cache for MITRE ATT&CK mapping.');
if (!fs.existsSync(storageDir)) {
// cache the newest version of ATT&CK
if (debug) { log('Cached directory does not exist. Creating.'); }
fs.mkdirSync(storageDir, {recursive: true});
result = await helpers.downloadLatestAttackMap(storageDir);
}
else {
// check the cache directory for matching files
const cachedPath: string|undefined = await helpers.getLatestCacheVersion(storageDir);
if (debug) { log(`Using cache path: ${cachedPath}`); }
if (cachedPath === undefined) {
// no files found - download the latest version from GitHub
log('Nothing found in extension cache. Downloading latest version of MITRE ATT&CK mapping');
Expand All @@ -66,17 +68,25 @@ export async function cacheData(storageDir: string): Promise<AttackMap|undefined
// files found - compare the cached version to the newest version on GitHub
// Example: enterprise-attack.8.0.json => 8.0
const cachedVersion = path.basename(cachedPath).replace('enterprise-attack.', '').replace('.json', '');
const availableVersions: Array<string> = await helpers.getVersions();
const onlineVersion = `${availableVersions.sort()[availableVersions.length - 1]}`;
if (cachedVersion < onlineVersion) {
// if online version is newer than the cached one, download and use the online version
vscode.window.showInformationMessage('ATT&CK: Identified a new version of the ATT&CK mapping! Replacing cached version.');
log(`Identified a new version of the ATT&CK mapping! Replacing cached map (${cachedVersion}) with downloaded map (${onlineVersion})`);
result = await helpers.downloadLatestAttackMap(storageDir);
}
else {
// otherwise just use the cached one
log(`Nothing to do. Cached version is on latest ATT&CK version ${onlineVersion}`);
if (debug) { log(`Cached version: ${cachedVersion}`); }
try {
const availableVersions: Array<string> = await helpers.getVersions();
const onlineVersion = `${availableVersions.sort()[availableVersions.length - 1]}`;
if (debug) { log(`Online version: ${onlineVersion}`); }
if (cachedVersion < onlineVersion) {
// if online version is newer than the cached one, download and use the online version
vscode.window.showInformationMessage('ATT&CK: Identified a new version of the ATT&CK mapping! Replacing cached version.');
log(`Identified a new version of the ATT&CK mapping! Replacing cached map (${cachedVersion}) with downloaded map (${onlineVersion})`);
result = await helpers.downloadLatestAttackMap(storageDir);
}
else {
// otherwise just use the cached one
log(`Nothing to do. Cached version is on latest ATT&CK version ${onlineVersion}`);
const cachedData: string = fs.readFileSync(cachedPath, {encoding: 'utf8'});
result = JSON.parse(cachedData) as AttackMap;
}
} catch (error) {
log(`Could not download ATT&CK version from GitHub. Falling back to cached version ${cachedVersion}.`);
const cachedData: string = fs.readFileSync(cachedPath, {encoding: 'utf8'});
result = JSON.parse(cachedData) as AttackMap;
}
Expand Down
47 changes: 38 additions & 9 deletions src/helpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@ export const mitigationRegex = /M\d{4}/;
// everything under this will only show the technique provider's results
export const minTermLength = 5;

const httpTimeout: number = 5000;

/*
Send a given message to the MITRE ATT&CK output channel with a timestamp
*/
Expand Down Expand Up @@ -69,9 +71,11 @@ export function getVersions(prefix = 'ATT&CK-v'): Promise<Array<string>> {
};
return new Promise((resolve, reject) => {
let filteredTags: Array<string> = new Array<string>();
https.get(url, options, (res: IncomingMessage) => {
const request = https.get(url, options, (res: IncomingMessage) => {
res.setEncoding('utf8');
res.on('data', (chunk) => { downloadedData = downloadedData.concat(chunk); });
res.on('data', (chunk) => {
downloadedData = downloadedData.concat(chunk);
});
res.on('error', (err: Error) => {
// something bad happened! let the user know
log(`Could not retrieve the version list! ${err.message}`);
Expand All @@ -94,8 +98,21 @@ export function getVersions(prefix = 'ATT&CK-v'): Promise<Array<string>> {
});
resolve(filteredTags);
}
else {
try {
const result: Record<string,string> = JSON.parse(downloadedData);
log(`Error encountered while downloading ATT&CK map versions: ${result['message']}`);
} catch (error) {
log(`No tags were parsed! Something went wrong! Is api.github.com reachable?`);
}
reject('ATT&CK map versions could not be downloaded');
}
});
});
request.setTimeout(httpTimeout, () => {
log(`HTTP request timed out while downloading ATT&CK map versions! Is api.github.com reachable?`);
reject('HTTP request timed out');
});
});
}

Expand Down Expand Up @@ -131,7 +148,7 @@ export function downloadAttackMap(storageDir: string, version: string): Promise<
// Example: v8.0 => enterprise-attack.8.0.json
const storagePath: string = path.join(storageDir, `enterprise-attack.${version}.json`);
const url = `https://raw.githubusercontent.com/mitre/cti/ATT%26CK-v${version}/enterprise-attack/enterprise-attack.json`;
https.get(url, (res: IncomingMessage) => {
const request = https.get(url, (res: IncomingMessage) => {
res.setEncoding('utf8');
res.on('data', (chunk) => { downloadedData = downloadedData.concat(chunk); });
res.on('error', (err: Error) => {
Expand All @@ -147,6 +164,10 @@ export function downloadAttackMap(storageDir: string, version: string): Promise<
resolve(downloadedData);
});
});
request.setTimeout(httpTimeout, () => {
log(`HTTP request timed out while downloading ATT&CK map ${version}! Is raw.githubusercontent.com reachable?`);
reject('HTTP request timed out');
});
}
else {
log(`Could not find version ${version} in the tags list: ${availableVersions}`);
Expand All @@ -161,12 +182,20 @@ export function downloadAttackMap(storageDir: string, version: string): Promise<
*/
export async function downloadLatestAttackMap(storageDir: string): Promise<AttackMap|undefined> {
let result: AttackMap|undefined = undefined;
const availableVersions: Array<string> = await getVersions();
// always look for the latest tagged version
const version = `${availableVersions.sort()[availableVersions.length - 1]}`;
const downloadedData: string = await downloadAttackMap(storageDir, version);
// and once it's cached, parse + return it
result = JSON.parse(downloadedData) as AttackMap;
try {
const availableVersions: Array<string> = await getVersions();
// always look for the latest tagged version
const version = `${availableVersions.sort()[availableVersions.length - 1]}`;
try {
const downloadedData: string = await downloadAttackMap(storageDir, version);
// and once it's cached, parse + return it
result = JSON.parse(downloadedData) as AttackMap;
} catch (err) {
log(`downloadLatestAttackMap() failed due to '${err}'`);
}
} catch (err) {
log(`getVersions() failed due to '${err}'`);
}
return result;
}

Expand Down
9 changes: 6 additions & 3 deletions src/test/suite/extension.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ describe('Extension', function () {

beforeEach(ignoreConsoleLogs);
afterEach(resetState);
it('should use a cached version of the ATT&CK technique data if available', function (done) {
it('should use a cached version of the ATT&CK map if available', function (done) {
const tmpDir: string = os.tmpdir();
const tmpPath: string = path.join(tmpDir, enterpriseAttackFilename);
// queue this up to be deleted after the test has finished
Expand All @@ -44,7 +44,7 @@ describe('Extension', function () {
});
});
});
it('should download a new version of the ATT&CK technique data if none is cached', async function () {
it('should download a new version of the ATT&CK map if none is cached', async function () {
const tmpDir: string = os.tmpdir();
// collect the current timestamp
const currTime: Date = new Date();
Expand All @@ -59,7 +59,7 @@ describe('Extension', function () {
// allow the cached file to be created within 15 seconds of current time, just for some wiggle room
assert.ok((currTime.getMilliseconds() - cachedFileStats.mtimeMs) < 15000);
});
it('should download a new version of the ATT&CK technique data if the available one is outdated', async function () {
it('should download a new version of the ATT&CK map if the available one is outdated', async function () {
const tmpDir: string = os.tmpdir();
const tmpPath: string = path.join(tmpDir, 'enterprise-attack.7.2.json');
const oldMapPath: string = path.join(__dirname, '..', '..', '..', 'src', 'test', 'files', 'attack7.json');
Expand All @@ -74,6 +74,9 @@ describe('Extension', function () {
assert.ok(attackMap !== undefined);
assert.strictEqual(helpers.isAttackMapNewer(attackMap, oldMap), true);
});
it.skip('should use a cached version of the ATT&CK map if the online version could not be downloaded', async function () {
// TODO
});
it.skip('toggleStatusBar: should display to the user when a matching file is active', async function () {
// TODO
});
Expand Down

0 comments on commit 5aa7a97

Please sign in to comment.