Introduce Trust Book in Book Viewlet (#9414)

extensions/notebook/src/book/bookTreeView.ts

@@ -7,19 +7,20 @@ import * as azdata from 'azdata';
 import * as vscode from 'vscode';
 import * as path from 'path';
 import * as fs from 'fs-extra';
+import * as constants from '../common/constants';
 import * as fsw from 'fs';
 import { IPrompter, QuestionTypes, IQuestion } from '../prompts/question';
 import CodeAdapter from '../prompts/adapter';
 import { BookTreeItem } from './bookTreeItem';
 import { BookModel } from './bookModel';
 import { Deferred } from '../common/promise';
+import { IBookTrustManager, BookTrustManager } from './bookTrustManager';
 import * as loc from '../common/localizedConstants';
 import { ApiWrapper } from '../common/apiWrapper';
 
 const Content = 'content';
 
 export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeItem> {
-
 	private _onDidChangeTreeData: vscode.EventEmitter<BookTreeItem | undefined> = new vscode.EventEmitter<BookTreeItem | undefined>();
 	readonly onDidChangeTreeData: vscode.Event<BookTreeItem | undefined> = this._onDidChangeTreeData.event;
 	private _throttleTimer: any;
@@ -27,21 +28,22 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 	private _extensionContext: vscode.ExtensionContext;
 	private prompter: IPrompter;
 	private _initializeDeferred: Deferred<void> = new Deferred<void>();
-	private _bookViewer: vscode.TreeView<BookTreeItem>;
 	private _openAsUntitled: boolean;
-	private _apiWrapper: ApiWrapper;
+	private _bookTrustManager: IBookTrustManager;
+
+	private _bookViewer: vscode.TreeView<BookTreeItem>;
 	public viewId: string;
 	public books: BookModel[];
 	public currentBook: BookModel;
 
-	constructor(apiWrapper: ApiWrapper, workspaceFolders: vscode.WorkspaceFolder[], extensionContext: vscode.ExtensionContext, openAsUntitled: boolean, view: string) {
+	constructor(private _apiWrapper: ApiWrapper, workspaceFolders: vscode.WorkspaceFolder[], extensionContext: vscode.ExtensionContext, openAsUntitled: boolean, view: string) {
 		this._openAsUntitled = openAsUntitled;
 		this._extensionContext = extensionContext;
 		this.books = [];
 		this.initialize(workspaceFolders).catch(e => console.error(e));
 		this.viewId = view;
 		this.prompter = new CodeAdapter();
-		this._apiWrapper = apiWrapper ? apiWrapper : new ApiWrapper();
+		this._bookTrustManager = new BookTrustManager(this.books, _apiWrapper);
 	}
 
 	private async initialize(workspaceFolders: vscode.WorkspaceFolder[]): Promise<void> {
@@ -60,6 +62,28 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 		return this._initializeDeferred.promise;
 	}
 
+	get _visitedNotebooks(): string[] {
+		return this._extensionContext.globalState.get(constants.visitedNotebooksMementoKey, []);
+	}
+
+	set _visitedNotebooks(value: string[]) {
+		this._extensionContext.globalState.update(constants.visitedNotebooksMementoKey, value);
+	}
+
+	trustBook(bookTreeItem?: BookTreeItem): void {
+		let bookPathToTrust = bookTreeItem ? bookTreeItem.root : this.currentBook?.bookPath;
+
+		if (bookPathToTrust) {
+			let trustChanged = this._bookTrustManager.setBookAsTrusted(bookPathToTrust);
+
+			if (trustChanged) {
+				this._apiWrapper.showInfoMessage(loc.msgBookTrusted);
+			} else {
+				this._apiWrapper.showInfoMessage(loc.msgBookAlreadyTrusted);
+			}
+		}
+	}
+
 	async openBook(bookPath: string, urlToOpen?: string): Promise<void> {
 		try {
 			let books: BookModel[] = this.books.filter(book => book.bookPath === bookPath) || [];
@@ -157,6 +181,19 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 			if (this._openAsUntitled) {
 				this.openNotebookAsUntitled(resource);
 			} else {
+				// let us keep a list of already visited notebooks so that we do not trust them again, potentially
+				// overriding user changes
+				let normalizedResource = path.normalize(resource);
+
+				if (this._visitedNotebooks.indexOf(normalizedResource) === -1
+					&& this._bookTrustManager.isNotebookTrustedByDefault(normalizedResource)) {
+					let openDocumentListenerUnsubscriber = azdata.nb.onDidOpenNotebookDocument((document: azdata.nb.NotebookDocument) => {
+						document.setTrusted(true);
+						this._visitedNotebooks = this._visitedNotebooks.concat([normalizedResource]);
+						openDocumentListenerUnsubscriber.dispose();
+					});
+				}
+
 				let doc = await vscode.workspace.openTextDocument(resource);
 				vscode.window.showTextDocument(doc);
 			}
@@ -377,6 +414,4 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 			default: false
 		});
 	}
-
-
 }

extensions/notebook/src/book/bookTrustManager.ts

@@ -0,0 +1,114 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+import * as path from 'path';
+import * as vscode from 'vscode';
+import * as constants from './../common/constants';
+import { BookTreeItem } from './bookTreeItem';
+import { BookModel } from './bookModel';
+import { ApiWrapper } from '../common/apiWrapper';
+
+export interface IBookTrustManager {
+	isNotebookTrustedByDefault(notebookUri: string): boolean;
+	setBookAsTrusted(bookRootPath: string): boolean;
+}
+
+enum TrustBookOperation {
+	Add,
+	Remove
+}
+
+export class BookTrustManager implements IBookTrustManager {
+	constructor(private books: BookModel[], private apiWrapper: ApiWrapper) { }
+
+	isNotebookTrustedByDefault(notebookUri: string): boolean {
+		let normalizedNotebookUri: string = path.normalize(notebookUri);
+		let treeBookItems: BookTreeItem[] = this.getBookTreeItems();
+		let trustableBookPaths = this.getTrustableBookPaths();
+		let hasTrustedBookPath: boolean = treeBookItems
+			.filter(bookItem => trustableBookPaths.some(trustableBookPath => trustableBookPath === path.join(bookItem.book.root, path.sep)))
+			.some(bookItem => normalizedNotebookUri.startsWith(path.join(bookItem.root, 'content', path.sep)));
+		let isNotebookTrusted = hasTrustedBookPath && this.books.some(bookModel => bookModel.getNotebook(normalizedNotebookUri));
+		return isNotebookTrusted;
+	}
+
+	getTrustableBookPaths() {
+		let trustablePaths: string[];
+		let bookPathsInConfig: string[] = this.getTrustedBookPathsInConfig();
+
+		if (this.hasWorkspaceFolders()) {
+			let workspaceFolders = this.apiWrapper.getWorkspaceFolders();
+
+			trustablePaths = bookPathsInConfig
+				.map(trustableBookPath => workspaceFolders
+					.map(workspaceFolder => path.join(workspaceFolder.uri.fsPath, trustableBookPath)))
+				.reduce((accumulator, currentTrustableBookPaths) => accumulator.concat(currentTrustableBookPaths), []);
+
+		} else {
+			trustablePaths = bookPathsInConfig;
+		}
+
+		return trustablePaths;
+	}
+
+	getBookTreeItems(): BookTreeItem[] {
+		return this.books
+			.map(book => book.bookItems) // select all the books
+			.reduce((accumulator, currentBookItemList) => accumulator.concat(currentBookItemList), []);
+	}
+
+	setBookAsTrusted(bookRootPath: string): boolean {
+		return this.updateTrustedBooks(bookRootPath, TrustBookOperation.Add);
+	}
+
+	getTrustedBookPathsInConfig(): string[] {
+		let config: vscode.WorkspaceConfiguration = this.apiWrapper.getConfiguration(constants.notebookConfigKey);
+		let trustedBookDirectories: string[] = config.get(constants.trustedBooksConfigKey);
+
+		return trustedBookDirectories;
+	}
+
+	setTrustedBookPathsInConfig(bookPaths: string[]) {
+		let config: vscode.WorkspaceConfiguration = this.apiWrapper.getConfiguration(constants.notebookConfigKey);
+		let storeInWorspace: boolean = this.hasWorkspaceFolders();
+
+		config.update(constants.trustedBooksConfigKey, bookPaths, storeInWorspace ? false : vscode.ConfigurationTarget.Global);
+	}
+
+	hasWorkspaceFolders(): boolean {
+		let workspaceFolders = this.apiWrapper.getWorkspaceFolders();
+		return workspaceFolders && workspaceFolders.length > 0;
+	}
+
+	updateTrustedBooks(bookPath: string, operation: TrustBookOperation) {
+		let modifiedTrustedBooks = false;
+		let bookPathToChange: string = path.join(bookPath, path.sep);
+
+		if (this.hasWorkspaceFolders()) {
+			let workspaceFolders = this.apiWrapper.getWorkspaceFolders();
+			let matchingWorkspaceFolder: vscode.WorkspaceFolder = workspaceFolders
+				.find(ws => bookPathToChange.startsWith(path.normalize(ws.uri.fsPath)));
+
+			// if notebook is stored in a workspace folder, then store only the relative directory
+			if (matchingWorkspaceFolder) {
+				bookPathToChange = bookPathToChange.replace(path.normalize(matchingWorkspaceFolder.uri.fsPath), '');
+			}
+		}
+
+		let trustedBooks: string[] = this.getTrustedBookPathsInConfig();
+		let existingBookIndex = trustedBooks.map(trustedBookPath => path.normalize(trustedBookPath)).indexOf(bookPathToChange);
+
+		if (existingBookIndex !== -1 && operation === TrustBookOperation.Remove) {
+			trustedBooks.splice(existingBookIndex, 1);
+			modifiedTrustedBooks = true;
+		} else if (existingBookIndex === -1 && operation === TrustBookOperation.Add) {
+			trustedBooks.push(bookPathToChange);
+			modifiedTrustedBooks = true;
+		}
+
+		this.setTrustedBookPathsInConfig(trustedBooks);
+
+		return modifiedTrustedBooks;
+	}
+}