Skip to content

Commit a71d36f

Browse files
stavinskistavinski
authored
Merge pull request #1 from stavinski/features
Features
2 parents b33d7d0 + c729242 commit a71d36f

1 file changed

Lines changed: 80 additions & 10 deletions

File tree

‎gui.py‎

Lines changed: 80 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1-
from javax.swing import JTabbedPane, JPanel, JButton, JLabel, SwingConstants, JOptionPane, GroupLayout, JCheckBox, JSplitPane
1+
from javax.swing import JTabbedPane, JPanel, JButton, JLabel, SwingConstants, JOptionPane, GroupLayout, JCheckBox, JSplitPane, JRadioButton, ButtonGroup, JFileChooser
22
from javax.swing.event import ChangeListener, DocumentListener
33
from javax.swing.LayoutStyle.ComponentPlacement import RELATED, UNRELATED
44
from java.awt import BorderLayout, Font, Component
55
from java.beans import PropertyChangeListener
66
from org.python.core.util import StringUtil
7+
from burp import IExtensionStateListener
78
from uicomponents import BurpUI, TabComponent, TabComponentEditableTabMixin, TabComponentCloseableMixin, TabComponentCloseListener, TabComponentTitleChangedListener
89
from models import ObservableCollection, Script
910
from utils import EditorFileAdapter
@@ -118,6 +119,7 @@ def __init__(self, callbacks, script):
118119
self.scriptEditor = callbacks.createTextEditor()
119120
self.scriptEditor.text = script.content
120121
self.scriptText = self.scriptEditor.component
122+
self.loadButton = JButton('Load', actionPerformed=self.load)
121123
self.compileButton = JButton('Compile', actionPerformed=self.compile, enabled=False)
122124

123125
editingLayout = GroupLayout(self, autoCreateGaps=True, autoCreateContainerGaps=True)
@@ -128,7 +130,10 @@ def __init__(self, callbacks, script):
128130
)
129131
.addGroup(editingLayout.createParallelGroup()
130132
.addComponent(self.scriptText)
131-
.addComponent(self.compileButton)
133+
)
134+
.addGroup(editingLayout.createSequentialGroup()
135+
.addComponent(self.loadButton)
136+
.addComponent(self.compileButton)
132137
)
133138
)
134139

@@ -137,8 +142,11 @@ def __init__(self, callbacks, script):
137142
.addComponent(self.enabledCheckbox)
138143
)
139144
.addGroup(editingLayout.createSequentialGroup()
140-
.addComponent(self.scriptText)
141-
.addComponent(self.compileButton)
145+
.addComponent(self.scriptText)
146+
)
147+
.addGroup(editingLayout.createParallelGroup()
148+
.addComponent(self.loadButton)
149+
.addComponent(self.compileButton)
142150
)
143151
)
144152
self.layout = editingLayout
@@ -148,6 +156,13 @@ def __init__(self, callbacks, script):
148156
def enabled_changed(self, event):
149157
self.script.enabled = self.enabledCheckbox.isSelected()
150158

159+
def load(self, event):
160+
file_chooser = JFileChooser()
161+
choice = file_chooser.showOpenDialog(None)
162+
if choice == JFileChooser.APPROVE_OPTION:
163+
with open(file_chooser.selectedFile.path, 'r') as input_file:
164+
self.scriptEditor.text = ''.join(input_file.readlines())
165+
151166
def compile(self, event):
152167
self.script.compile()
153168
self.compileButton.enabled = False
@@ -173,8 +188,8 @@ def _can_compile(self, event):
173188
self.compileButton.enabled = self.script.requires_compile
174189

175190

176-
class ScriptOutputPanel(JPanel, PropertyChangeListener):
177-
191+
class ScriptOutputPanel(JPanel, PropertyChangeListener, IExtensionStateListener):
192+
178193
def __init__(self, callbacks, script):
179194
super(ScriptOutputPanel, self).__init__()
180195
self.callbacks = callbacks
@@ -189,37 +204,93 @@ def __init__(self, callbacks, script):
189204
self.add(self.tabbedPane, BorderLayout.CENTER)
190205
self.script.stdout = EditorFileAdapter(self.outputEditor)
191206
self.script.stderr = EditorFileAdapter(self.errorEditor)
207+
self.output_file = None
208+
209+
# register to be notified when the extension is unloaded so if a output_file ref is in use it can be closed
210+
callbacks.registerExtensionStateListener(self)
192211

193212
def clear_stderr(self, event):
194213
self.errorEditor.text = ''
195214

196215
def clear_stdout(self, event):
197216
self.outputEditor.text = ''
198217

218+
def save_file_output(self, event):
219+
self.outputFileBrowseButton.enabled = True
220+
if self.output_file: # already have an output_file
221+
self.script.stdout = self.output_file
222+
return
223+
224+
if not self.set_output_file():
225+
# didn't choose a file then revert to using UI
226+
self.outputUIRadioButton.selected = True
227+
228+
def view_ui_output(self, event):
229+
self.outputFileBrowseButton.enabled = False
230+
self.script.stdout = EditorFileAdapter(self.outputEditor)
231+
232+
def set_output_file(self, event=None):
233+
file_chooser = JFileChooser()
234+
choice = file_chooser.showSaveDialog(None)
235+
if choice == JFileChooser.APPROVE_OPTION:
236+
self.outputFileLabel.text = file_chooser.selectedFile.path
237+
self.output_file = open(file_chooser.selectedFile.path, 'a')
238+
self.script.stdout = self.output_file
239+
return True
240+
241+
return False
242+
199243
def propertyChange(self, event):
200244
if event.propertyName == Script.Properties.IS_COMPILED:
201245
if event.newValue:
202246
self.errorEditor.text = ''
203247
elif event.propertyName == Script.Properties.COMPILATION_ERROR:
204248
self.errorEditor.text = event.newValue
205-
self.tabbedPane.selectedIndex = 1
249+
self.tabbedPane.selectedIndex = 1
250+
251+
def extensionUnloaded(self):
252+
if self.output_file: # if we have a file ref then close it
253+
print('Closing output file reference')
254+
self.output_file.close()
206255

207256
def _create_output_panel(self):
208257
self.outputPanel = JPanel()
209258
self.outputEditor = self.callbacks.createTextEditor()
210259
self.outputEditor.editable = False
211260
self.outputText = self.outputEditor.component
212261
self.clearOutputButton = JButton('Clear', actionPerformed=self.clear_stdout)
262+
self.outputButtonGroup = ButtonGroup()
263+
self.outputFileRadioButton = JRadioButton('Save to File:', actionPerformed=self.save_file_output)
264+
self.outputUIRadioButton = JRadioButton('Show in UI:', selected=True, actionPerformed=self.view_ui_output)
265+
self.outputFileLabel = JLabel()
266+
self.outputFileBrowseButton = JButton('Browse...', enabled=False, actionPerformed=self.set_output_file)
267+
268+
self.outputButtonGroup.add(self.outputFileRadioButton)
269+
self.outputButtonGroup.add(self.outputUIRadioButton)
213270

214271
outputLayout = GroupLayout(self.outputPanel, autoCreateGaps=True, autoCreateContainerGaps=True)
215272
outputLayout.setHorizontalGroup(outputLayout.createParallelGroup()
273+
.addGroup(
274+
outputLayout.createSequentialGroup()
275+
.addComponent(self.outputFileRadioButton)
276+
.addComponent(self.outputFileLabel)
277+
.addComponent(self.outputFileBrowseButton)
278+
)
279+
.addComponent(self.outputUIRadioButton)
216280
.addComponent(self.outputText)
217281
.addComponent(self.clearOutputButton)
218282
)
219283

220284
outputLayout.setVerticalGroup(outputLayout.createSequentialGroup()
221-
.addComponent(self.outputText)
222-
.addComponent(self.clearOutputButton)
285+
.addGroup(
286+
outputLayout.createParallelGroup()
287+
.addComponent(self.outputFileRadioButton)
288+
.addComponent(self.outputFileLabel)
289+
.addComponent(self.outputFileBrowseButton)
290+
)
291+
.addComponent(self.outputUIRadioButton)
292+
.addComponent(self.outputText)
293+
.addComponent(self.clearOutputButton)
223294
)
224295
self.outputPanel.layout = outputLayout
225296

@@ -241,7 +312,6 @@ def _create_error_panel(self):
241312
)
242313
self.errorPanel.layout = errorLayout
243314

244-
245315
class ScriptPanel(JPanel):
246316

247317
def __init__(self, script, callbacks):

0 commit comments

Comments
 (0)