diff --git a/bin/threatdown.ts b/bin/threatdown.ts index 4e1b253..e95e85b 100755 --- a/bin/threatdown.ts +++ b/bin/threatdown.ts @@ -1,18 +1,16 @@ #!/usr/bin/env node import { readFileSync, writeFileSync } from "node:fs"; -import { resolve } from "node:path"; +import { extname, resolve } from "node:path"; import { parseArgs } from "node:util"; import { parse, compileToMermaid, + generateUpdatedMd, renderMermaid, } from "../lib"; -const { - values, - positionals, -} = parseArgs({ +const { values, positionals } = parseArgs({ allowPositionals: true, options: { output: { @@ -27,22 +25,48 @@ const { }, }); -function usage () { +function usage() { console.log(`Usage: threatdown + Must have extension \`.td\` or \`.md\` --output Write result to file --type Change output type, must be one of "json", "mermaid" or "svg" `); } -async function main () { - const inputFile = positionals.shift(); +async function main() { + if (positionals.length !== 1) { + return usage(); + } + + // non-null assertion safe because we already checked for length + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + const inputFile = positionals.shift()!; + const inputFileExt = extname(inputFile); + // non-null assertion safe because the outputType has a default // eslint-disable-next-line @typescript-eslint/no-non-null-assertion - if (!inputFile || !["json", "mermaid", "svg"].includes(values.type!)) { - usage(); - } else { - const fileContent = readFileSync(resolve(process.cwd(), inputFile), { encoding: "utf8" }); + if ( + !values.type || + !["json", "mermaid", "svg"].includes(values.type) || + ![".md", ".td"].includes(inputFileExt) + ) { + process.exitCode = 1; + return usage(); + } + + const fileContent = readFileSync(resolve(process.cwd(), inputFile), { + encoding: "utf8", + }); + + if (inputFileExt === ".md") { + const markdownContent = await generateUpdatedMd(fileContent, values.type); + if (values.output) { + writeFileSync(resolve(process.cwd(), values.output), markdownContent); + } else { + console.log(markdownContent); + } + } else if (inputFileExt === ".td") { const parsedContent = parse(fileContent); if (values.type === "json") { if (values.output) { @@ -50,6 +74,7 @@ async function main () { } else { console.log(JSON.stringify(parsedContent, null, 2)); } + return; } @@ -60,6 +85,7 @@ async function main () { } else { console.log(mermaidContent); } + return; } @@ -71,11 +97,12 @@ async function main () { console.log(svgContent); } } + + return; } } -main() - .catch((err: Error) => { - process.exitCode = 1; - console.error(err.stack); - }); +main().catch((err: Error) => { + process.exitCode = 1; + console.error(err.stack); +}); diff --git a/lib/index.ts b/lib/index.ts index cb1608b..26a0f4e 100644 --- a/lib/index.ts +++ b/lib/index.ts @@ -1,3 +1,4 @@ export { parse } from "./parser"; export { compileToMermaid } from "./compiler"; export { renderMermaid } from "./renderer"; +export { generateUpdatedMd } from "./md-converter"; diff --git a/lib/md-converter.ts b/lib/md-converter.ts new file mode 100644 index 0000000..a07ccc2 --- /dev/null +++ b/lib/md-converter.ts @@ -0,0 +1,78 @@ +import { parse } from "./parser/index"; +import { compileToMermaid } from "./compiler"; +import { renderMermaid } from "./renderer"; + +const threatdownRegex = /```threatdown([\s\S]*?)```/g; + +const getMermaidSvg = async (mermaidFormatedData: string): Promise => { + const testSvg = await renderMermaid(mermaidFormatedData); + return testSvg; +}; + +const cleanThreatdownBlocks = (input: string) => { + return input + .trim() + .replace(/^```threatdown/, "") + .replace(/```$/, ""); +}; + +// Define an asynchronous function to process each match +const processMatchAsync = async ( + match: string, + mdEmbeddedType: string +): Promise => { + const cleanMatch = cleanThreatdownBlocks(match); + const jsonFormatedData = parse(cleanMatch); + const mermaidRaw = compileToMermaid(jsonFormatedData); + let mermaidSvg = ""; + if (mdEmbeddedType === "svg") { + mermaidSvg = await getMermaidSvg(mermaidRaw); + } + + return ( + ` \n` + + // Render Mermaid + `${ + mdEmbeddedType === "json" + ? "```json\n" + JSON.stringify(jsonFormatedData) + "\n```\n" + : "" + }` + + // Render Mermaid + `${ + mdEmbeddedType === "mermaid" + ? "```mermaid\n" + mermaidRaw + "\n```\n" + : "" + }` + + // Render SVG + `${mdEmbeddedType === "svg" ? "\n" + mermaidSvg + "\n" : ""}` + ); +}; + +// generate the updated markdown file +export const generateUpdatedMd = async ( + file: string, + mdEmbeddedType: string +): Promise => { + try { + const threatdownMatches = file.match(threatdownRegex); + + if (!threatdownMatches) { + throw new Error("No threatdown content found"); + } + + const newFilePromises = threatdownMatches.map((match) => + processMatchAsync(match, mdEmbeddedType) + ); + const newFileContentArray = await Promise.all(newFilePromises); + + let newFile = file; + for (let i = 0; i < threatdownMatches.length; i++) { + newFile = newFile.replace(threatdownMatches[i], newFileContentArray[i]); + } + + return newFile; + } catch (err) { + console.error(err); + throw err; + } +}; diff --git a/test/fixtures/sample.md b/test/fixtures/sample.md new file mode 100644 index 0000000..6065f6e --- /dev/null +++ b/test/fixtures/sample.md @@ -0,0 +1,18 @@ +## This is a markdown document + +```threatdown +__Attacker's goal__ + - method which in order to be viable + + (high) requires this condition to be true + + and this condition which depends on either + - x to be true + - or y to be true + + hey this condition must be true too + - another method here too + - a condition which depends on assumptions + +? this might be a problem + +? but only if this happens + -? which assumes this also happens +``` + +### Above this is a graph