├── requirements.txt ├── requirements-3.8.txt ├── images ├── disk.png ├── printer.png ├── zenmap.png ├── add-node.png ├── edit-bold.png ├── edit-list.png ├── question.png ├── scissors.png ├── RedNotebook.png ├── arrow-curve.png ├── disk--pencil.png ├── edit-color.png ├── edit-italic.png ├── nodes │ ├── folder.png │ ├── os_win.png │ ├── os_apple.png │ ├── os_linux.png │ ├── question.png │ ├── stat_red.png │ ├── os_freebsd.png │ ├── stat_green.png │ ├── stat_yellow.png │ ├── checkbox_checked.png │ ├── checkbox_tasks.png │ └── checkbox_unchecked.png ├── ui-tab--plus.png ├── add-root-node.png ├── arrow-continue.png ├── document-copy.png ├── edit-alignment.png ├── edit-underline.png ├── edit-list-order.png ├── selection-input.png ├── arrow-curve-180-left.png ├── edit-alignment-center.png ├── edit-alignment-justify.png ├── edit-alignment-right.png ├── blue-folder-open-document.png └── clipboard-paste-document-text.png ├── base.py ├── init.sql ├── catalog.py ├── LICENSE ├── Readme.md └── redteamnotebook.py /requirements.txt: -------------------------------------------------------------------------------- 1 | libnmap 2 | PyQt5>=5.6 3 | sip 4 | sqlalchemy 5 | -------------------------------------------------------------------------------- /requirements-3.8.txt: -------------------------------------------------------------------------------- 1 | natlas-libnmap 2 | PyQt5>=5.6 3 | sip 4 | sqlalchemy 5 | -------------------------------------------------------------------------------- /images/disk.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/disk.png -------------------------------------------------------------------------------- /base.py: -------------------------------------------------------------------------------- 1 | from sqlalchemy.ext.declarative import declarative_base 2 | Base = declarative_base() 3 | -------------------------------------------------------------------------------- /images/printer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/printer.png -------------------------------------------------------------------------------- /images/zenmap.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/zenmap.png -------------------------------------------------------------------------------- /images/add-node.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/add-node.png -------------------------------------------------------------------------------- /images/edit-bold.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/edit-bold.png -------------------------------------------------------------------------------- /images/edit-list.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/edit-list.png -------------------------------------------------------------------------------- /images/question.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/question.png -------------------------------------------------------------------------------- /images/scissors.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/scissors.png -------------------------------------------------------------------------------- /images/RedNotebook.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/RedNotebook.png -------------------------------------------------------------------------------- /images/arrow-curve.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/arrow-curve.png -------------------------------------------------------------------------------- /images/disk--pencil.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/disk--pencil.png -------------------------------------------------------------------------------- /images/edit-color.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/edit-color.png -------------------------------------------------------------------------------- /images/edit-italic.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/edit-italic.png -------------------------------------------------------------------------------- /images/nodes/folder.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/folder.png -------------------------------------------------------------------------------- /images/nodes/os_win.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/os_win.png -------------------------------------------------------------------------------- /images/ui-tab--plus.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/ui-tab--plus.png -------------------------------------------------------------------------------- /images/add-root-node.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/add-root-node.png -------------------------------------------------------------------------------- /images/arrow-continue.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/arrow-continue.png -------------------------------------------------------------------------------- /images/document-copy.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/document-copy.png -------------------------------------------------------------------------------- /images/edit-alignment.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/edit-alignment.png -------------------------------------------------------------------------------- /images/edit-underline.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/edit-underline.png -------------------------------------------------------------------------------- /images/nodes/os_apple.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/os_apple.png -------------------------------------------------------------------------------- /images/nodes/os_linux.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/os_linux.png -------------------------------------------------------------------------------- /images/nodes/question.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/question.png -------------------------------------------------------------------------------- /images/nodes/stat_red.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/stat_red.png -------------------------------------------------------------------------------- /images/edit-list-order.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/edit-list-order.png -------------------------------------------------------------------------------- /images/nodes/os_freebsd.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/os_freebsd.png -------------------------------------------------------------------------------- /images/nodes/stat_green.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/stat_green.png -------------------------------------------------------------------------------- /images/nodes/stat_yellow.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/stat_yellow.png -------------------------------------------------------------------------------- /images/selection-input.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/selection-input.png -------------------------------------------------------------------------------- /images/arrow-curve-180-left.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/arrow-curve-180-left.png -------------------------------------------------------------------------------- /images/edit-alignment-center.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/edit-alignment-center.png -------------------------------------------------------------------------------- /images/edit-alignment-justify.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/edit-alignment-justify.png -------------------------------------------------------------------------------- /images/edit-alignment-right.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/edit-alignment-right.png -------------------------------------------------------------------------------- /images/nodes/checkbox_checked.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/checkbox_checked.png -------------------------------------------------------------------------------- /images/nodes/checkbox_tasks.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/checkbox_tasks.png -------------------------------------------------------------------------------- /images/nodes/checkbox_unchecked.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/nodes/checkbox_unchecked.png -------------------------------------------------------------------------------- /images/blue-folder-open-document.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/blue-folder-open-document.png -------------------------------------------------------------------------------- /images/clipboard-paste-document-text.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unix-ninja/redteamnotebook/HEAD/images/clipboard-paste-document-text.png -------------------------------------------------------------------------------- /init.sql: -------------------------------------------------------------------------------- 1 | CREATE TABLE node_graph ( 2 | nodeid TEXT, 3 | parentid TEXT, 4 | basename TEXT, 5 | icon TEXT, 6 | mtime FLOAT, 7 | UNIQUE(nodeid) 8 | ); 9 | 10 | CREATE TABLE notes ( 11 | nodeid TEXT, 12 | content TEXT, 13 | mtime FLOAT, 14 | UNIQUE(nodeid) ON CONFLICT REPLACE 15 | ); 16 | -------------------------------------------------------------------------------- /catalog.py: -------------------------------------------------------------------------------- 1 | from sqlalchemy import Column, String, Integer, Float, Boolean, ForeignKey 2 | from sqlalchemy.orm import relationship, backref 3 | from sqlalchemy.ext.declarative import declarative_base 4 | from base import Base 5 | 6 | class NodeGraph(Base): 7 | __tablename__ = "node_graph" 8 | nodeid = Column(String, primary_key=True) 9 | parentid = Column(String) 10 | basename = Column(String) 11 | icon = Column(String) 12 | mtime = Column(Float) 13 | 14 | class Note(Base): 15 | __tablename__ = "notes" 16 | nodeid = Column(String, ForeignKey("node_graph.nodeid", ondelete="CASCADE"), primary_key=True) 17 | content = Column(String) 18 | mtime = Column(Float) 19 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 2-Clause License 2 | 3 | Copyright (c) 2020, unix-ninja 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | * Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | * Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 17 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 18 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 19 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 20 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 21 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 22 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 23 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 24 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | -------------------------------------------------------------------------------- /Readme.md: -------------------------------------------------------------------------------- 1 | # Readteam Notebook 2 | 3 | Redteam Notebook is an experiment to address digital intelligence archiving on offensive engagements. It should be easy to take notes and screenshots, organize data, and import security tooling output from other common sources. 4 | 5 | **NOTE:** This is currently alpha, and is not feature complete, but should work as an MVP. 6 | 7 | ## Setting it up 8 | 9 | It is recommended to use a Python virtual environment (though it's not a requirement to do so.) You can follow the steps below. 10 | 11 | ### Using pipenv 12 | 13 | Make sure you have pipenv installed. You can install with pip3 or your package manager before proceeding. 14 | 15 | ``` 16 | $ git clone https://github.com/unix-ninja/redteamnotebook.git 17 | $ cd redteamnotebook/ 18 | $ pipenv --python 3 19 | $ pipenv install -r requirements.txt 20 | $ pipenv run python redteamnotebook.py 21 | ``` 22 | 23 | ### Using python3 venv 24 | 25 | Instead of a third party package like pipenv, you could try venv which is now built into Python 3. No additional setup is required. 26 | 27 | ``` 28 | $ git clone https://github.com/unix-ninja/redteamnotebook.git 29 | $ cd redteamnotebook/ 30 | $ python3 -m venv . 31 | $ source bin/activate 32 | $ pip install -r requirements.txt 33 | $ python redteamnotebook.py 34 | ``` 35 | 36 | > Note: For Python 3.8, please use the alternate requirements-3.8.txt file. This replaces libnmap with natlas-libnmap, as the former is not currently available for Python 3.8: 37 | 38 | ``` 39 | $ pipenv install -r requirements-3.8.txt 40 | ``` 41 | 42 | ## Quick Start 43 | 44 | Once you have launched Redteam Notebook, you should have a new, empty notebook up on your screen. The first thing you will want to do is add a new root node. Click on the "New Root Node" button in the toolbar. You should see a new node called 'Node' appear in the left pane. Click on that node, and now you can start placing notes in the right-hand pane. 45 | 46 | You can add children nodes by clicking on a node and then using the 'New Node' button, or additional root nodes with the 'New Root Node' button again. 47 | 48 | For more information, visit https://www.unix-ninja.com/p/introducing_redteam_notebook 49 | -------------------------------------------------------------------------------- /redteamnotebook.py: -------------------------------------------------------------------------------- 1 | # TODO: error checks for insert nodes (text + parentid should be unique) 2 | # TODO: better markdown editing 3 | from PyQt5.QtGui import * 4 | from PyQt5.QtWidgets import * 5 | from PyQt5.QtCore import * 6 | from PyQt5.QtPrintSupport import * 7 | 8 | import argparse 9 | import platform 10 | import sqlalchemy 11 | import subprocess 12 | import catalog 13 | 14 | from libnmap.parser import NmapParser 15 | import hashlib 16 | import json 17 | import os 18 | import shutil 19 | import sys 20 | import uuid 21 | 22 | from math import ceil 23 | 24 | APP_PATH = os.path.dirname(os.path.realpath(__file__))+'/' 25 | TEXT_STYLES = ['Title', 'Heading', 'Subheading', 'Body'] 26 | TEXT_LEVEL = {"Body": 0, "Title": 1, "Heading": 2, "Subheading": 3} 27 | TEXT_SIZE = {"Body": 14, "Title": 26, "Heading": 20, "Subheading": 14} 28 | TEXT_WEIGHT = {"Body": 50, "Title": 75, "Heading": 75, "Subheading": 75} 29 | IMAGE_EXTENSIONS = ['.jpg','.jpeg','.png','.bmp'] 30 | HTML_EXTENSIONS = ['.htm', '.html'] 31 | NODE_ICON_PATH = os.path.abspath(APP_PATH+'/images/nodes') 32 | ROLE_NODE_UUID = Qt.UserRole + 1 33 | NOTEBOOK_PATH = os.path.abspath(os.path.expanduser('~/default.notebook')) 34 | SETTINGS = os.path.abspath(os.path.expanduser('~/.local/redteamnotebook.cfg')) 35 | OS_ICONS = {'Windows': 'os_win.png', 'Linux': 'os_linux.png', 'Mac OS X': 'os_apple.png', 'FreeBSD': 'os_freebsd.png' } 36 | 37 | ## 38 | settings = { 39 | 'last_open_notebook': NOTEBOOK_PATH 40 | } 41 | 42 | ## initialize Session for the db 43 | Session = None 44 | 45 | 46 | def hexuuid(): 47 | return uuid.uuid4().hex 48 | 49 | def splitext(p): 50 | return os.path.splitext(p)[1].lower() 51 | 52 | def info(text, level=None): 53 | global args 54 | 55 | map = { 56 | 'debug': '[debug] ', 57 | 'error': '[err] ', 58 | 'info': '[info] ' 59 | } 60 | prefix='' 61 | 62 | if level == 'debug' and not args.debug: 63 | return 64 | if level in map: 65 | prefix=map[level] 66 | print(prefix+text) 67 | 68 | def move_node(uuid=None, parentid=None): 69 | db = Session() 70 | node = db.query(catalog.NodeGraph).get(uuid) 71 | ## update the icon in the catalog 72 | if node: 73 | node.parentid = parentid 74 | db.add(node) 75 | db.commit() 76 | else: 77 | info ('Unable to find node in catalog.', level='error') 78 | 79 | class StandardItem(QStandardItem): 80 | def __init__(self, txt='', font_size=14, fullref=None, uuid=None, set_bold=False, color=QColor(0, 0, 0)): 81 | super().__init__() 82 | 83 | fnt = QFont('Helvetica', font_size) 84 | fnt.setBold(set_bold) 85 | 86 | self.setEditable(True) 87 | self.setForeground(color) 88 | self.setFont(fnt) 89 | self.setText(txt) 90 | self.setToolTip(txt) 91 | self.setIcon(QIcon(os.path.join(NODE_ICON_PATH, 'folder.png'))) 92 | self.setEditable(True) 93 | self.setData(fullref, Qt.UserRole) 94 | 95 | if not uuid: 96 | uuid = hexuuid() 97 | self.setData(uuid, ROLE_NODE_UUID) 98 | 99 | class CAction(QWidgetAction): 100 | colorSelected = pyqtSignal(QColor) 101 | 102 | def __init__(self, parent): 103 | QWidgetAction.__init__(self, parent) 104 | widget = QWidget(parent) 105 | layout = QGridLayout(widget) 106 | layout.setSpacing(0) 107 | layout.setContentsMargins(2, 2, 2, 2) 108 | 109 | ## grab icons from our node path 110 | icons = os.listdir(NODE_ICON_PATH) 111 | icons.sort(reverse=True) 112 | count = len(icons) 113 | rows = count // round(count ** .5) 114 | for row in range(rows): 115 | for column in range(5): 116 | if not len(icons): break 117 | icon = "" 118 | while len(icons) and not icon.endswith('.png'): 119 | icon = icons.pop() 120 | button = QToolButton(widget) 121 | button.setAutoRaise(True) 122 | button.clicked.connect(lambda : self.handleButton(button)) 123 | button.setIcon(QIcon(os.path.join(NODE_ICON_PATH, icon))) 124 | button.setText(icon) 125 | layout.addWidget(button, row, column) 126 | self.setDefaultWidget(widget) 127 | 128 | def handleButton(self, button): 129 | ## close the context menu 130 | self.parent().hide() 131 | 132 | ## grab the item 133 | window = self.parent().parent() 134 | treeView = window.treeView 135 | treeModel = window.treeModel 136 | idx = treeView.selectedIndexes() 137 | if not idx: return 138 | item = treeModel.itemFromIndex(idx[0]) 139 | uuid = item.data(ROLE_NODE_UUID) 140 | 141 | ## change the icon 142 | item.setIcon(self.sender().icon()) 143 | 144 | ## fetch the node from the catalog 145 | db = Session() 146 | node = db.query(catalog.NodeGraph).get(uuid) 147 | ## update the icon in the catalog 148 | if node: 149 | node.icon = self.sender().text() 150 | db.add(node) 151 | db.commit() 152 | else: 153 | info ('Unable to find node in catalog.', level='error') 154 | 155 | class CMenu(QMenu): 156 | def __init__(self, parent): 157 | QMenu.__init__(self, parent) 158 | self.colorAction = CAction(self) 159 | self.colorAction.colorSelected.connect(self.handleColorSelected) 160 | self.addAction(self.colorAction) 161 | self.addSeparator() 162 | 163 | def handleColorSelected(self, color): 164 | print(color.name()) 165 | 166 | class TextEdit(QTextEdit): 167 | def canInsertFromMimeData(self, source): 168 | 169 | if source.hasImage(): 170 | return True 171 | else: 172 | return super(TextEdit, self).canInsertFromMimeData(source) 173 | 174 | def insertFromMimeData(self, source): 175 | 176 | cursor = self.textCursor() 177 | document = self.document() 178 | max_width = self.size().width() 179 | staging_file = 'images/stage.png' 180 | images = [] 181 | 182 | if source.hasUrls(): 183 | ## set image from file 184 | for u in source.urls(): 185 | file_ext = splitext(str(u.toLocalFile())) 186 | if u.isLocalFile() and file_ext in IMAGE_EXTENSIONS: 187 | images.append(QImage(u.toLocalFile())) 188 | elif source.hasImage(): 189 | ## set image from clipboard content 190 | images.append(source.imageData()) 191 | 192 | ## save and add our image if we have one 193 | if images: 194 | for image in images: 195 | ## make sure we insert images on blank lines. Insert one if we need to 196 | cursor.movePosition(QTextCursor.StartOfBlock) 197 | cursor.movePosition(QTextCursor.EndOfBlock, QTextCursor.KeepAnchor) 198 | if cursor.selectedText().strip(): 199 | cursor.movePosition(QTextCursor.EndOfBlock) 200 | cursor.insertText("\n") 201 | 202 | ## save our staging file 203 | image.save(os.path.abspath(staging_file)) 204 | ## get a file hash to rename the staging file 205 | img_hash = hashlib.md5(open(staging_file,'rb').read()).hexdigest() 206 | ## save the file in the notebook 207 | saved_file = f"images/{img_hash}.png" 208 | shutil.move(staging_file, saved_file) 209 | ## resize img if it's larger than the editor 210 | if image.width() > max_width: 211 | image = image.scaledToWidth(max_width) 212 | ## add image to the doc 213 | cursor.insertImage(saved_file) 214 | block = cursor.block() 215 | ## we want to add a newline after our last image if it's the end of the document 216 | if cursor.blockNumber() == document.blockCount() - 1: 217 | cursor.insertText("\n") 218 | ## make sure we resize our images after inserting 219 | self.resizeImages() 220 | ## when we finish processing our images, just return 221 | return 222 | 223 | ## If we hit a non-image or non-local URL, fall out to the super call & let Qt handle it 224 | super(TextEdit, self).insertFromMimeData(source) 225 | 226 | def resizeImages(self): 227 | document = self.document() 228 | max_width = self.size().width() - 10 229 | cursor = self.textCursor() 230 | #cursor.setPosition(0) 231 | 232 | block = document.begin() 233 | while block != document.end(): 234 | it = block.begin() 235 | while not it.atEnd(): 236 | fragment = it.fragment() 237 | if fragment.isValid(): 238 | if fragment.charFormat().isImageFormat(): 239 | img_fmt = fragment.charFormat().toImageFormat() 240 | ## let's figure out our max image size 241 | image = QImage() 242 | image.load(img_fmt.name()) 243 | new_width = max_width if image.width() > max_width else image.width() 244 | img_fmt.setWidth(new_width) 245 | cursor.setPosition(fragment.position()) 246 | cursor.setPosition(fragment.position() + fragment.length(), QTextCursor.KeepAnchor) 247 | cursor.setCharFormat(img_fmt) 248 | it += 1 249 | block = block.next() 250 | 251 | return 252 | 253 | def keyPressEvent(self, event): 254 | if (event.key() == Qt.Key_Return): 255 | ## we want to insert a regular formatted block when enter is pressed 256 | cursor = self.textCursor() 257 | blockFormat = QTextBlockFormat() 258 | blockFormat.setHeadingLevel(0) 259 | charFormat = QTextCharFormat() 260 | charFormat.setFontPointSize(14) 261 | cursor.insertBlock(blockFormat, charFormat) 262 | elif (event.key() == Qt.Key_Escape): 263 | cursor = self.textCursor() 264 | block = cursor.blockNumber() 265 | bf = cursor.blockFormat() 266 | cf = cursor.charFormat() 267 | 268 | print ('block: '+ str(block)) 269 | print ('weight: '+ str(cf.fontWeight())) 270 | self.resizeImages() 271 | else: 272 | QTextEdit.keyPressEvent(self, event) 273 | 274 | def mouseDoubleClickEvent(self,event): 275 | super().mouseDoubleClickEvent(event) 276 | cursor = self.textCursor() 277 | block = cursor.block() 278 | it = block.begin() 279 | while not it.atEnd(): 280 | fragment = it.fragment() 281 | if fragment.isValid(): 282 | if fragment.charFormat().isImageFormat(): 283 | img_fmt = fragment.charFormat().toImageFormat() 284 | self.openfile(img_fmt.name()) 285 | break 286 | it += 1 287 | 288 | def openfile(self, filename): 289 | if platform.system() == 'Darwin': # macOS 290 | subprocess.call(('open', filename)) 291 | elif platform.system() == 'Windows': # Windows 292 | os.startfile(filename) 293 | else: # linux variants 294 | subprocess.call(('xdg-open', filename)) 295 | 296 | def onContentsChanged(self): 297 | if self.updating: return 298 | self.save_doc = True 299 | self.updating = True 300 | 301 | ## we can do some format checking here later if we want to 302 | 303 | self.updating = False 304 | 305 | return 306 | 307 | def set_style(self, style): 308 | if style not in TEXT_STYLES: 309 | return 310 | 311 | ## set our cursor 312 | cursor = self.textCursor() 313 | cursor.movePosition(QTextCursor.StartOfBlock) 314 | cursor.movePosition(QTextCursor.EndOfBlock, QTextCursor.KeepAnchor) 315 | 316 | ## set our formatting 317 | blockFormat = QTextBlockFormat() 318 | blockFormat.setHeadingLevel(TEXT_LEVEL[style]) 319 | charFormat = QTextCharFormat() 320 | charFormat.setFontPointSize(TEXT_SIZE[style]) 321 | charFormat.setFontWeight(TEXT_WEIGHT[style]) 322 | cursor.setBlockFormat(blockFormat) 323 | cursor.setCharFormat(charFormat) 324 | 325 | class CDialog(QDialog): 326 | 327 | def __init__(self, *args, **kwargs): 328 | super(CDialog, self).__init__(*args, **kwargs) 329 | 330 | self.setWindowTitle("Add port") 331 | 332 | QBtn = QDialogButtonBox.Ok | QDialogButtonBox.Cancel 333 | 334 | self.buttonBox = QDialogButtonBox(QBtn) 335 | self.buttonBox.accepted.connect(self.accept) 336 | self.buttonBox.rejected.connect(self.reject) 337 | 338 | self.port = QLineEdit(self) 339 | self.proto = QComboBox(self) 340 | self.state = QComboBox(self) 341 | 342 | ## set our protocols 343 | self.proto.addItem('tcp') 344 | self.proto.addItem('udp') 345 | 346 | ## set our states 347 | self.state.addItem('open') 348 | self.state.addItem('closed') 349 | self.state.addItem('filtered') 350 | 351 | ## set up our layout 352 | self.layout = QFormLayout() 353 | self.layout.addRow("Port", self.port) 354 | self.layout.addRow("Protocol", self.proto) 355 | self.layout.addRow("State", self.state) 356 | 357 | self.layout.addWidget(self.buttonBox) 358 | self.setLayout(self.layout) 359 | 360 | class CTreeView(QTreeView): 361 | def __init__(self): 362 | QTreeView.__init__(self) 363 | self.setDragEnabled(True) 364 | self.setAcceptDrops(True) 365 | self.setDropIndicatorShown(True) 366 | self.setDragDropMode(QAbstractItemView.InternalMove) 367 | self.setUniformRowHeights(True) 368 | 369 | def startDrag(self, actions): 370 | self._node = self.selectedIndexes()[0] 371 | self._prev_parent = self._node.parent() 372 | return QTreeView.startDrag(self, actions) 373 | 374 | def dropEvent(self, event): 375 | idx = self.indexAt(event.pos()) 376 | dip = self.dropIndicatorPosition() 377 | if dip == QAbstractItemView.AboveItem: 378 | parent = idx.parent() 379 | elif dip == QAbstractItemView.BelowItem: 380 | parent = idx.parent() 381 | elif dip == QAbstractItemView.OnItem: 382 | parent = idx 383 | elif dip == QAbstractItemView.OnViewport: 384 | pass 385 | 386 | if not parent: 387 | parent = self.treeModel.invisibleRootItem() 388 | 389 | uuid = self._node.data(ROLE_NODE_UUID) 390 | parentid = parent.data(ROLE_NODE_UUID) 391 | 392 | move_node(uuid, parentid) 393 | 394 | super().dropEvent(event) 395 | 396 | 397 | class MainWindow(QMainWindow): 398 | def __init__(self, *args, **kwargs): 399 | super(MainWindow, self).__init__(*args, **kwargs) 400 | 401 | layout = QGridLayout() 402 | layout.setColumnStretch(1,1) 403 | layout.setSpacing(0) 404 | layout.setContentsMargins(0,0,0,0) 405 | 406 | self.docs = {} 407 | self.editor = TextEdit() 408 | self.editor.updating = False 409 | self.editor.new_line = False 410 | ## Setup the QTextEdit editor configuration 411 | self.editor.setAutoFormatting(QTextEdit.AutoAll) 412 | self.editor.selectionChanged.connect(self.update_format) 413 | self.editor.cursorPositionChanged.connect(self.monitor_style) 414 | ## Initialize default font size. 415 | font = QFont('Helvetica', 14) 416 | self.editor.setFont(font) 417 | ## We need to repeat the size to init the current format. 418 | self.editor.setFontPointSize(14) 419 | ## start editor disabled, until we click a node 420 | self.editor.setReadOnly(True) 421 | 422 | ## self.path holds the path of the currently open file. 423 | ## If none, we haven't got a file open yet (or creating new). 424 | self.path = None 425 | 426 | self.treeView = CTreeView() 427 | self.treeView.setStyleSheet("QTreeView { selection-background-color: #c3e3ff;} ") 428 | 429 | self.treeModel = QStandardItemModel() 430 | self.treeModel.setHorizontalHeaderLabels(['Targets']) 431 | rootNode = self.treeModel.invisibleRootItem() 432 | 433 | ## populate our tree 434 | self.load_nodes_from_catalog(clean=True) 435 | 436 | ## use the model with our view 437 | self.treeView.setModel(self.treeModel) 438 | self.treeView.clicked.connect(self.fetch_note) 439 | self.treeView.setContextMenuPolicy(Qt.CustomContextMenu) 440 | self.treeView.customContextMenuRequested.connect(self.show_context_menu) 441 | 442 | layout.addWidget(self.treeView,1,0) 443 | layout.addWidget(self.editor,1,1) 444 | 445 | self.treeModel.dataChanged.connect(self.tree_changed) 446 | container = QWidget() 447 | container.setLayout(layout) 448 | self.setCentralWidget(container) 449 | 450 | self.status = QStatusBar() 451 | self.setStatusBar(self.status) 452 | 453 | file_toolbar = QToolBar("File") 454 | file_toolbar.setIconSize(QSize(14, 14)) 455 | self.addToolBar(file_toolbar) 456 | file_menu = self.menuBar().addMenu("&File") 457 | 458 | open_file_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'blue-folder-open-document.png')), "Open notebook...", self) 459 | open_file_action.setStatusTip("Open notebook") 460 | open_file_action.triggered.connect(self.file_open) 461 | file_menu.addAction(open_file_action) 462 | file_toolbar.addAction(open_file_action) 463 | 464 | new_root_node_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'add-root-node.png')), "New Root Node", self) 465 | new_root_node_action.setStatusTip("New Root Node") 466 | new_root_node_action.triggered.connect(self.add_root_node) 467 | file_menu.addAction(new_root_node_action) 468 | file_toolbar.addAction(new_root_node_action) 469 | 470 | new_node_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'add-node.png')), "New Node", self) 471 | new_node_action.setStatusTip("New Node") 472 | new_node_action.triggered.connect(self.add_node) 473 | file_menu.addAction(new_node_action) 474 | file_toolbar.addAction(new_node_action) 475 | 476 | import_nmap_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'zenmap.png')), "Import NMap", self) 477 | import_nmap_action.setStatusTip("Import NMap") 478 | import_nmap_action.triggered.connect(self.import_nmap) 479 | file_menu.addAction(import_nmap_action) 480 | file_toolbar.addAction(import_nmap_action) 481 | 482 | print_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'printer.png')), "Print...", self) 483 | print_action.setStatusTip("Print current page") 484 | print_action.triggered.connect(self.file_print) 485 | file_menu.addAction(print_action) 486 | file_toolbar.addAction(print_action) 487 | 488 | edit_toolbar = QToolBar("Edit") 489 | edit_toolbar.setIconSize(QSize(16, 16)) 490 | self.addToolBar(edit_toolbar) 491 | edit_menu = self.menuBar().addMenu("&Edit") 492 | 493 | undo_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'arrow-curve-180-left.png')), "Undo", self) 494 | undo_action.setStatusTip("Undo last change") 495 | undo_action.triggered.connect(self.editor.undo) 496 | edit_menu.addAction(undo_action) 497 | 498 | redo_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'arrow-curve.png')), "Redo", self) 499 | redo_action.setStatusTip("Redo last change") 500 | redo_action.triggered.connect(self.editor.redo) 501 | edit_toolbar.addAction(redo_action) 502 | edit_menu.addAction(redo_action) 503 | 504 | edit_menu.addSeparator() 505 | 506 | cut_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'scissors.png')), "Cut", self) 507 | cut_action.setStatusTip("Cut selected text") 508 | cut_action.setShortcut(QKeySequence.Cut) 509 | cut_action.triggered.connect(self.editor.cut) 510 | edit_toolbar.addAction(cut_action) 511 | edit_menu.addAction(cut_action) 512 | 513 | copy_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'document-copy.png')), "Copy", self) 514 | copy_action.setStatusTip("Copy selected text") 515 | cut_action.setShortcut(QKeySequence.Copy) 516 | copy_action.triggered.connect(self.editor.copy) 517 | edit_toolbar.addAction(copy_action) 518 | edit_menu.addAction(copy_action) 519 | 520 | paste_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'clipboard-paste-document-text.png')), "Paste", self) 521 | paste_action.setStatusTip("Paste from clipboard") 522 | cut_action.setShortcut(QKeySequence.Paste) 523 | paste_action.triggered.connect(self.editor.paste) 524 | edit_toolbar.addAction(paste_action) 525 | edit_menu.addAction(paste_action) 526 | 527 | select_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'selection-input.png')), "Select all", self) 528 | select_action.setStatusTip("Select all text") 529 | cut_action.setShortcut(QKeySequence.SelectAll) 530 | select_action.triggered.connect(self.editor.selectAll) 531 | edit_menu.addAction(select_action) 532 | 533 | edit_menu.addSeparator() 534 | 535 | wrap_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'arrow-continue.png')), "Wrap text to window", self) 536 | wrap_action.setStatusTip("Toggle wrap text to window") 537 | wrap_action.setCheckable(True) 538 | wrap_action.setChecked(True) 539 | wrap_action.triggered.connect(self.edit_toggle_wrap) 540 | edit_menu.addAction(wrap_action) 541 | 542 | format_toolbar = QToolBar("Format") 543 | format_toolbar.setIconSize(QSize(16, 16)) 544 | self.addToolBar(format_toolbar) 545 | format_menu = self.menuBar().addMenu("&Format") 546 | 547 | ## We need references to these actions/settings to update as selection changes, so attach to self. 548 | self.fonts = QFontComboBox() 549 | self.fonts.currentFontChanged.connect(self.editor.setCurrentFont) 550 | 551 | self.stylebox = QComboBox() 552 | self.stylebox.addItems(TEXT_STYLES) 553 | 554 | self.stylebox.setCurrentIndex(3) 555 | self.stylebox.currentIndexChanged[str].connect(lambda s: self.setStyle(s) ) 556 | format_toolbar.addWidget(self.stylebox) 557 | 558 | self.bold_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'edit-bold.png')), "Bold", self) 559 | self.bold_action.setStatusTip("Bold") 560 | self.bold_action.setShortcut(QKeySequence.Bold) 561 | self.bold_action.setCheckable(True) 562 | self.bold_action.toggled.connect(lambda x: self.editor.setFontWeight(QFont.Bold if x else QFont.Normal)) 563 | format_toolbar.addAction(self.bold_action) 564 | format_menu.addAction(self.bold_action) 565 | 566 | self.italic_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'edit-italic.png')), "Italic", self) 567 | self.italic_action.setStatusTip("Italic") 568 | self.italic_action.setShortcut(QKeySequence.Italic) 569 | self.italic_action.setCheckable(True) 570 | self.italic_action.toggled.connect(self.editor.setFontItalic) 571 | format_toolbar.addAction(self.italic_action) 572 | format_menu.addAction(self.italic_action) 573 | 574 | self.underline_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'edit-underline.png')), "Underline", self) 575 | self.underline_action.setStatusTip("Underline") 576 | self.underline_action.setShortcut(QKeySequence.Underline) 577 | self.underline_action.setCheckable(True) 578 | self.underline_action.toggled.connect(self.editor.setFontUnderline) 579 | format_toolbar.addAction(self.underline_action) 580 | format_menu.addAction(self.underline_action) 581 | 582 | format_menu.addSeparator() 583 | 584 | self.alignl_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'edit-alignment.png')), "Align left", self) 585 | self.alignl_action.setStatusTip("Align text left") 586 | self.alignl_action.setCheckable(True) 587 | self.alignl_action.triggered.connect(lambda: self.editor.setAlignment(Qt.AlignLeft)) 588 | 589 | self.alignc_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'edit-alignment-center.png')), "Align center", self) 590 | self.alignc_action.setStatusTip("Align text center") 591 | self.alignc_action.setCheckable(True) 592 | self.alignc_action.triggered.connect(lambda: self.editor.setAlignment(Qt.AlignCenter)) 593 | 594 | self.alignr_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'edit-alignment-right.png')), "Align right", self) 595 | self.alignr_action.setStatusTip("Align text right") 596 | self.alignr_action.setCheckable(True) 597 | self.alignr_action.triggered.connect(lambda: self.editor.setAlignment(Qt.AlignRight)) 598 | 599 | self.alignj_action = QAction(QIcon(os.path.join(APP_PATH+'images', 'edit-alignment-justify.png')), "Justify", self) 600 | self.alignj_action.setStatusTip("Justify text") 601 | self.alignj_action.setCheckable(True) 602 | self.alignj_action.triggered.connect(lambda: self.editor.setAlignment(Qt.AlignJustify)) 603 | 604 | format_group = QActionGroup(self) 605 | format_group.setExclusive(True) 606 | format_group.addAction(self.alignl_action) 607 | format_group.addAction(self.alignc_action) 608 | format_group.addAction(self.alignr_action) 609 | format_group.addAction(self.alignj_action) 610 | 611 | format_menu.addSeparator() 612 | 613 | ## A list of all format-related widgets/actions, so we can disable/enable signals when updating. 614 | self._format_actions = [ 615 | self.fonts, 616 | self.bold_action, 617 | self.italic_action, 618 | self.underline_action, 619 | ## We don't need to disable signals for alignment, as they are paragraph-wide. 620 | ] 621 | 622 | ## change toolbar styles 623 | toolbar_style = "windows" 624 | file_toolbar.setMovable(False) 625 | edit_toolbar.setMovable(False) 626 | self.status.setStyle(QStyleFactory.create(toolbar_style)) 627 | 628 | ## setup a resize timer here. We'll use this to resize images when the main window resizes 629 | self.resize_timer = QTimer() 630 | 631 | ## Initialize. 632 | self.update_format() 633 | self.update_title() 634 | self.resize(700,400) 635 | self.show() 636 | 637 | ## set flag for auto saves 638 | self.editor.save_doc = False 639 | 640 | ## make sure our timer isn't running from first resize. Then connect it. 641 | self.resize_timer.stop() 642 | self.resize_timer.timeout.connect(self.editor.resizeImages) 643 | 644 | ## setup our timer to auto save docs 645 | timer = QTimer(self) 646 | timer.timeout.connect(self.timeout_save) 647 | timer.start(5000) 648 | 649 | self.installEventFilter(self) 650 | 651 | def resizeEvent(self, event): 652 | self.resize_timer.stop() 653 | self.resize_timer.start(250) 654 | super().resizeEvent(event) 655 | 656 | def monitor_style(self): 657 | if self.editor.updating == True: 658 | return 659 | ## set our cursor 660 | cursor = self.editor.textCursor() 661 | blockFormat = cursor.blockFormat() 662 | 663 | ## lock updates before we change the style 664 | self.editor.updating = True 665 | ## make sure our style box reflects the style under the cursor 666 | for level in TEXT_LEVEL: 667 | if (blockFormat.headingLevel() == TEXT_LEVEL[level]): 668 | self.stylebox.setCurrentText(level) 669 | self.editor.updating = False 670 | 671 | def setStyle(self, style): 672 | self.editor.set_style(style) 673 | 674 | def timeout_save(self): 675 | if self.editor.save_doc: 676 | ## save editor content to catalog 677 | db = Session() 678 | note = catalog.Note() 679 | note.nodeid = self.get_nodeid() 680 | note.content = self.editor.toMarkdown() 681 | ## only save if the nodeid is valid 682 | if (not note.nodeid): 683 | self.editor.save_doc = False 684 | return 685 | db.add(note) 686 | db.commit() 687 | info ("Saved.", level='debug') 688 | self.editor.save_doc = False 689 | 690 | def tree_changed(self, signal): 691 | ## see what changed 692 | root = self.treeView.model().invisibleRootItem() 693 | node = self.treeView.selectedIndexes() 694 | if not node: 695 | info ('Unable to find selectedIndexes().', level='error') 696 | return 697 | basename = node[0].data(Qt.DisplayRole) 698 | uuid = node[0].data(ROLE_NODE_UUID) 699 | ## fetch the node from the catalog 700 | db = Session() 701 | node = db.query(catalog.NodeGraph).get(uuid) 702 | ## update the basename in the catalog 703 | if node: 704 | node.basename = basename 705 | db.add(node) 706 | db.commit() 707 | else: 708 | info (f'Unable to find node ({uuid}) in catalog.', level='error') 709 | return 710 | 711 | def show_context_menu(self, position): 712 | ## do we have a selection? 713 | node = self.treeView.selectedIndexes() 714 | 715 | ## build our menu 716 | menu = CMenu(self) 717 | 718 | if node: 719 | new_port_action = QAction("&Add Port") 720 | new_port_action.triggered.connect(self.add_port) 721 | menu.addAction(new_port_action) 722 | 723 | delete_node_action = QAction("&Remove Node") 724 | delete_node_action.triggered.connect(self.delete_node) 725 | menu.addAction(delete_node_action) 726 | 727 | ## show our menu 728 | menu.exec_(self.sender().viewport().mapToGlobal(position)) 729 | 730 | 731 | def add_port(self): 732 | node = self.treeView.selectedIndexes()[0] 733 | if not node: return None 734 | 735 | ## grab our id for later 736 | parentid = node.data(ROLE_NODE_UUID) 737 | 738 | dlg = CDialog(self) 739 | dlg.setWindowTitle("Add a port") 740 | accepted = dlg.exec_() 741 | if not accepted: return 742 | 743 | port = dlg.port.text() 744 | proto = dlg.proto.currentText() 745 | state = dlg.state.currentText() 746 | 747 | ## init our node vars 748 | proto_node = None 749 | # TODO: map desc to service name 750 | desc = '' 751 | 752 | if node.data(Qt.DisplayRole) == proto: 753 | proto_node = self.treeModel.itemFromIndex(node) 754 | else: 755 | ## find our protocol node 756 | rootNode = self.treeModel.itemFromIndex(node) 757 | for item in self.iterItems(rootNode): 758 | if item.data(Qt.DisplayRole) == proto: 759 | proto_node = item 760 | 761 | if not proto_node: 762 | ## add a node if we couldn't find one 763 | uuid = self.add_node(name=proto, parentid=parentid) 764 | else: 765 | uuid = proto_node.data(ROLE_NODE_UUID) 766 | 767 | ## set our node icon 768 | if state == 'closed': 769 | icon = 'stat_red.png' 770 | elif state == 'filtered': 771 | icon = 'stat_yellow.png' 772 | else: 773 | icon = 'stat_green.png' 774 | 775 | ## add port to protocol node 776 | portid = self.add_node(name=f'{port} {desc} [{state}]', parentid=uuid, icon=icon) 777 | 778 | def delete_node(self): 779 | node = self.treeView.selectedIndexes()[0] 780 | if not node: return None 781 | confirmed = QMessageBox.question(self, "Delete", f"Are you sure you want to delete '{node.data(Qt.DisplayRole)}'?", QMessageBox.Yes|QMessageBox.No) 782 | if confirmed == QMessageBox.No: 783 | return 784 | 785 | node = self.treeView.selectedIndexes()[0] 786 | if not node: return None 787 | 788 | db = Session() 789 | rootNode = self.treeModel.itemFromIndex(node) 790 | 791 | ## remove children from catalog 792 | for item in self.iterItems(rootNode): 793 | ## remove node graph 794 | db_node = db.query(catalog.NodeGraph).get(item.data(ROLE_NODE_UUID)) 795 | if db_node: 796 | db.delete(db_node) 797 | ## remove note 798 | db_node = db.query(catalog.Note).get(item.data(ROLE_NODE_UUID)) 799 | if db_node: 800 | db.delete(db_node) 801 | 802 | 803 | ## remove node from catalog 804 | db_node = db.query(catalog.NodeGraph).get(rootNode.data(ROLE_NODE_UUID)) 805 | if db_node: 806 | db.delete(db_node) 807 | db_node = db.query(catalog.Note).get(rootNode.data(ROLE_NODE_UUID)) 808 | if db_node: 809 | db.delete(db_node) 810 | 811 | db.commit() 812 | 813 | ## remove node from tree 814 | self.treeModel.removeRow(node.row(), parent=node.parent()) 815 | 816 | def iterItems(self, root): 817 | def recurse(parent): 818 | for row in range(parent.rowCount()): 819 | for column in range(parent.columnCount()): 820 | child = parent.child(row, column) 821 | yield child 822 | if child.hasChildren(): 823 | yield from recurse(child) 824 | if root is not None: 825 | yield from recurse(root) 826 | 827 | def load_nodes_from_catalog(self, parentid=None, clean=False): 828 | ## if clean is set, clear out tree and docs before loading catalog 829 | if clean: 830 | self.docs = {} 831 | rootNode = self.treeModel.invisibleRootItem() 832 | if (rootNode.hasChildren()): 833 | rootNode.removeRows(0, rootNode.rowCount()) 834 | 835 | ## load data from sql 836 | db = Session() 837 | nodes = db.query(catalog.NodeGraph).filter_by(parentid=parentid).all() 838 | for node in nodes: 839 | if not parentid: 840 | self.add_root_node(name=node.basename, uuid=node.nodeid, icon=node.icon) 841 | else: 842 | self.add_node(name=node.basename, uuid=node.nodeid, parentid=parentid, icon=node.icon) 843 | ## add child nodes, recursively 844 | self.load_nodes_from_catalog(node.nodeid) 845 | ## set content of this node 846 | note = db.query(catalog.Note).get(node.nodeid) 847 | if note: 848 | self.docs[node.nodeid].setMarkdown(note.content) 849 | return 850 | 851 | def get_nodeid(self): 852 | node = self.treeView.selectedIndexes() 853 | if not node: return None 854 | fullref = node[0].data(Qt.UserRole) 855 | uuid = node[0].data(ROLE_NODE_UUID) 856 | return uuid 857 | 858 | def itemFromUUID(self, uuid): 859 | rootNode = self.treeModel.invisibleRootItem() 860 | for item in self.iterItems(rootNode): 861 | if item.data(ROLE_NODE_UUID) == uuid: 862 | return item 863 | return None 864 | 865 | def fetch_note(self, signal): 866 | uuid = self.get_nodeid() 867 | 868 | if not uuid: return 869 | 870 | ## make sure we don't write the text we just loaded 871 | self.editor.updating = True 872 | ## display the proper doc 873 | self.editor.setDocument(self.docs[uuid]) 874 | ## make sure we can edit 875 | self.editor.setReadOnly(False) 876 | ## resize images on initial load 877 | self.editor.resizeImages() 878 | ## allow saving changes again 879 | self.editor.updating = False 880 | 881 | def add_root_node(self, name='Node', uuid=None, icon=None): 882 | record_catalog = False 883 | if not name: 884 | name = 'Node' 885 | if not uuid: 886 | uuid = hexuuid() 887 | record_catalog = True 888 | info ('Recording in catalog...', level='info') 889 | rootNode = self.treeModel.invisibleRootItem() 890 | fullref = '/'+name 891 | new_node = StandardItem(name, 14, fullref=fullref, uuid=uuid) 892 | if icon: 893 | new_node.setIcon(QIcon(os.path.join(NODE_ICON_PATH, icon))) 894 | rootNode.appendRow(new_node) 895 | idx = self.itemFromUUID(uuid) 896 | ## select the new node in the tree 897 | if idx: 898 | self.treeView.setCurrentIndex(idx.index()) 899 | 900 | ## create a doc on this node and allow it to be saved 901 | self.docs[uuid] = QTextDocument() 902 | doc = self.docs[uuid] 903 | doc.contentsChange.connect(self.editor.onContentsChanged) 904 | 905 | if record_catalog: 906 | ## record in catalog 907 | db = Session() 908 | node = catalog.NodeGraph() 909 | node.nodeid = uuid 910 | node.parentid = None 911 | node.basename = name 912 | db.add(node) 913 | db.commit() 914 | 915 | def add_node(self, name='Node', uuid=None, parentid=None, icon=None): 916 | record_catalog = False 917 | if not name: 918 | name = 'Node' 919 | if not uuid: 920 | uuid = hexuuid() 921 | record_catalog = True 922 | 923 | ## we will either be given a parent id or check for the selected item in the tree 924 | idx = None 925 | if parentid: 926 | rootNode = self.treeModel.invisibleRootItem() 927 | for item in self.iterItems(rootNode): 928 | if item.data(ROLE_NODE_UUID) == parentid: 929 | parent_node = item 930 | break 931 | else: 932 | idx = self.treeView.selectedIndexes() 933 | if not idx: return 934 | parent_node = self.treeModel.itemFromIndex(idx[0]) 935 | 936 | parent_fullref = parent_node.data(Qt.UserRole) 937 | fullref = f'{parent_fullref}/Node' 938 | 939 | new_node = StandardItem(name, 14, fullref=fullref, uuid=uuid) 940 | if icon: 941 | new_node.setIcon(QIcon(os.path.join(NODE_ICON_PATH, icon))) 942 | parent_node.appendRow(new_node) 943 | 944 | ## create a doc on this node and allow it to be saved 945 | self.docs[uuid] = QTextDocument() 946 | doc = self.docs[uuid] 947 | doc.contentsChange.connect(self.editor.onContentsChanged) 948 | 949 | if idx: 950 | self.treeView.setExpanded(idx[0], True) 951 | 952 | if record_catalog: 953 | info ('Recording in catalog...', level='info') 954 | ## record in catalog 955 | db = Session() 956 | node = catalog.NodeGraph() 957 | node.nodeid = uuid 958 | node.parentid = parent_node.data(ROLE_NODE_UUID) 959 | node.basename = name 960 | node.icon = icon 961 | db.add(node) 962 | db.commit() 963 | 964 | return uuid 965 | 966 | def block_signals(self, objects, b): 967 | for o in objects: 968 | o.blockSignals(b) 969 | 970 | def update_format(self): 971 | """ 972 | Update the font format toolbar/actions when a new text selection is made. This is neccessary to keep 973 | toolbars/etc. in sync with the current edit state. 974 | :return: 975 | """ 976 | ## Disable signals for all format widgets, so changing values here does not trigger further formatting. 977 | self.block_signals(self._format_actions, True) 978 | 979 | self.fonts.setCurrentFont(self.editor.currentFont()) 980 | 981 | self.italic_action.setChecked(self.editor.fontItalic()) 982 | self.underline_action.setChecked(self.editor.fontUnderline()) 983 | self.bold_action.setChecked(self.editor.fontWeight() == QFont.Bold) 984 | 985 | self.alignl_action.setChecked(self.editor.alignment() == Qt.AlignLeft) 986 | self.alignc_action.setChecked(self.editor.alignment() == Qt.AlignCenter) 987 | self.alignr_action.setChecked(self.editor.alignment() == Qt.AlignRight) 988 | self.alignj_action.setChecked(self.editor.alignment() == Qt.AlignJustify) 989 | 990 | self.block_signals(self._format_actions, False) 991 | 992 | def dialog_critical(self, s): 993 | dlg = QMessageBox(self) 994 | dlg.setText(s) 995 | dlg.setIcon(QMessageBox.Critical) 996 | dlg.show() 997 | 998 | def file_open(self): 999 | global NOTEBOOK_PATH 1000 | 1001 | ## lock updates 1002 | self.save_doc = False 1003 | self.updating = True 1004 | 1005 | dialog = QFileDialog() 1006 | dialog.setFileMode(QFileDialog.DirectoryOnly) 1007 | dialog.exec() 1008 | path = dialog.selectedFiles() 1009 | 1010 | if not path: 1011 | return 1012 | 1013 | ## don't reopen the same notebook 1014 | new_path = os.path.abspath(os.path.expanduser(path[0])) 1015 | if NOTEBOOK_PATH == new_path: 1016 | return 1017 | 1018 | ## we should init the notebook here 1019 | NOTEBOOK_PATH = new_path 1020 | info (f'Opening notebook "{NOTEBOOK_PATH}"', level='info') 1021 | 1022 | ## change the session to match the new file 1023 | set_session() 1024 | 1025 | ## init the notebook 1026 | init_notebook() 1027 | 1028 | self.load_nodes_from_catalog(clean=True) 1029 | self.update_title() 1030 | self.editor.setDocument(None) 1031 | self.editor.setReadOnly(True) 1032 | 1033 | ## update configs 1034 | settings['last_open_notebook'] = NOTEBOOK_PATH 1035 | save_settings() 1036 | 1037 | ## unlock 1038 | self.updating = False 1039 | 1040 | def file_print(self): 1041 | dlg = QPrintDialog() 1042 | if dlg.exec_(): 1043 | self.editor.print_(dlg.printer()) 1044 | 1045 | def update_title(self): 1046 | self.setWindowTitle("%s - Redteam Notebook" % (os.path.basename(self.path) if self.path else "Untitled")) 1047 | 1048 | def edit_toggle_wrap(self): 1049 | self.editor.setLineWrapMode( 1 if self.editor.lineWrapMode() == 0 else 0 ) 1050 | 1051 | def import_nmap(self): 1052 | msg = QMessageBox() 1053 | idx = self.treeView.selectedIndexes() 1054 | if not idx: 1055 | msg.setIcon(QMessageBox.Warning) 1056 | msg.setText("Please select a node before importing.") 1057 | msg.setStandardButtons(QMessageBox.Ok) 1058 | msg.exec_() 1059 | return 1060 | 1061 | ## make sure we've selected the first index 1062 | idx = idx[0] 1063 | 1064 | ## grab our id for later 1065 | parentid = idx.data(ROLE_NODE_UUID) 1066 | 1067 | ## open a dialog to select our file 1068 | dialog = QFileDialog() 1069 | filter = 'nmap xml file (*.xml)' 1070 | filename = dialog.getOpenFileName(None, 'Import NMap XML', '', filter)[0] 1071 | 1072 | ## If we cancelled the dialog, just return 1073 | if not filename: 1074 | return 1075 | 1076 | ## make sure the filename is valid 1077 | if not os.path.exists(filename): 1078 | msg.setIcon(QMessageBox.Critical) 1079 | msg.setText("Unable to open file!") 1080 | msg.setStandardButtons(QMessageBox.Ok) 1081 | msg.exec_() 1082 | return 1083 | 1084 | ## lock updates 1085 | self.save_doc = False 1086 | self.updating = True 1087 | 1088 | ## read xml file 1089 | nmap_report = NmapParser.parse_fromfile(filename) 1090 | 1091 | ## load results into tree 1092 | for host in nmap_report.hosts: 1093 | if not host.is_up(): 1094 | continue 1095 | hostid = None 1096 | icon = 'question.png' 1097 | ## generate the host label 1098 | if host.hostnames: 1099 | label = f'{host.address} ({host.hostnames[0]})' 1100 | else: 1101 | label = f'{host.address}' 1102 | 1103 | ## search for an OS icon 1104 | for c in host.os_class_probabilities(): 1105 | if c.osfamily in OS_ICONS: 1106 | icon = OS_ICONS[c.osfamily] 1107 | break 1108 | 1109 | ## add our node 1110 | hostid = self.add_node(name=label, parentid=parentid, icon=icon) 1111 | if not hostid: continue 1112 | 1113 | for service in host.services: 1114 | ## make sure we put services in the correct node 1115 | ### find our protocol node 1116 | proto_node = None 1117 | rootNode = self.itemFromUUID(hostid) 1118 | for item in self.iterItems(rootNode): 1119 | if item.data(Qt.DisplayRole) == service.protocol: 1120 | proto_node = item 1121 | 1122 | if not proto_node: 1123 | ### add a node if we couldn't find one 1124 | uuid = self.add_node(name=service.protocol, parentid=rootNode.data(ROLE_NODE_UUID)) 1125 | else: 1126 | ### use the node we found 1127 | uuid = proto_node.data(ROLE_NODE_UUID) 1128 | 1129 | ## set our node icon 1130 | if service.state == 'closed': 1131 | icon = 'stat_red.png' 1132 | elif service.state == 'filtered': 1133 | icon = 'stat_yellow.png' 1134 | else: 1135 | icon = 'stat_green.png' 1136 | ## add port to protocol node 1137 | portid = self.add_node(name=f'{service.port} {service.protocol} [{service.state}]', parentid=uuid, icon=icon) 1138 | 1139 | self.updating = False 1140 | 1141 | ## END MAIN WINDOW CLASS 1142 | 1143 | def init_sql(sql_path): 1144 | global args 1145 | info ('Setting up sql...', level='info') 1146 | ## create our tables 1147 | db_engine = sqlalchemy.create_engine(f'sqlite:///{NOTEBOOK_PATH}/catalog.sqlite', convert_unicode=True, echo=args.debug) 1148 | 1149 | ## apparently, sqlalchemy does not yet support ON CASCADE REPLACE, so we need to pass 1150 | ## some raw SQL to create our schema 1151 | 1152 | db = db_engine.connect() 1153 | 1154 | sql = """CREATE TABLE IF NOT EXISTS node_graph ( 1155 | nodeid TEXT, 1156 | parentid TEXT, 1157 | basename TEXT, 1158 | icon TEXT, 1159 | mtime FLOAT, 1160 | UNIQUE(nodeid) 1161 | );""" 1162 | result = db.execute(sql) 1163 | 1164 | sql = """CREATE TABLE IF NOT EXISTS notes ( 1165 | nodeid TEXT, 1166 | content TEXT, 1167 | mtime FLOAT, 1168 | UNIQUE(nodeid) ON CONFLICT REPLACE, 1169 | CONSTRAINT fk_nodeid 1170 | FOREIGN KEY (nodeid) 1171 | REFERENCES node_graph(nodeid) 1172 | ON DELETE CASCADE 1173 | );""" 1174 | result = db.execute(sql) 1175 | 1176 | def set_session(): 1177 | global Session 1178 | db_engine = sqlalchemy.create_engine(f'sqlite:///{NOTEBOOK_PATH}/catalog.sqlite', convert_unicode=True) 1179 | Session = sqlalchemy.orm.sessionmaker(bind=db_engine) 1180 | 1181 | def init_notebook(): 1182 | ## create the default notebook if it doesn't exist 1183 | if not os.path.exists(NOTEBOOK_PATH): 1184 | os.mkdir(NOTEBOOK_PATH) 1185 | if not os.path.exists(NOTEBOOK_PATH): 1186 | info ('Unable to create default notebook.', level='error') 1187 | sys.exit(1) 1188 | if not os.path.exists(NOTEBOOK_PATH+'/images'): 1189 | os.mkdir(NOTEBOOK_PATH+'/images') 1190 | if not os.path.exists(NOTEBOOK_PATH+'/catalog.sqlite'): 1191 | init_sql(NOTEBOOK_PATH) 1192 | if not Session: 1193 | set_session() 1194 | ## make sure our cwd is the path of the notebook 1195 | os.chdir(NOTEBOOK_PATH) 1196 | 1197 | def save_settings(): 1198 | with open(SETTINGS, 'w') as fp: 1199 | json.dump(settings, fp) 1200 | 1201 | if __name__ == '__main__': 1202 | ## parse arguments 1203 | parser = argparse.ArgumentParser(description='Redteam Notebook') 1204 | parser.add_argument('--debug', dest='debug', action='store_true', help='enable debug messages') 1205 | args = parser.parse_args() 1206 | 1207 | ## load settings 1208 | if not os.path.exists(SETTINGS): 1209 | save_settings() 1210 | else: 1211 | with open(SETTINGS) as fp: 1212 | settings = json.load(fp) 1213 | NOTEBOOK_PATH = settings['last_open_notebook'] 1214 | 1215 | ## get our configs 1216 | init_notebook() 1217 | 1218 | app = QApplication(sys.argv) 1219 | app.setApplicationName("Redteam Notebook") 1220 | 1221 | window = MainWindow() 1222 | app.exec_() 1223 | --------------------------------------------------------------------------------