From 0a256b3889c15a36ba7966c3597112b96cbb9346 Mon Sep 17 00:00:00 2001 From: Matteo Oldani Date: Thu, 4 May 2023 23:16:44 +0200 Subject: [PATCH 1/3] Made header scraper work on-demand. The history is now parsed to extract domain specific headers once the user opens the custom headers panel. --- python/inql/extender.py | 4 +- python/inql/scanner/customheaders.py | 88 ++++++++++------- python/inql/scraper/headers_scraper.py | 130 +++++++++++++++++++++++++ 3 files changed, 188 insertions(+), 34 deletions(-) diff --git a/python/inql/extender.py b/python/inql/extender.py index 26e70b4..f3e3c98 100644 --- a/python/inql/extender.py +++ b/python/inql/extender.py @@ -97,7 +97,9 @@ def registerExtenderCallbacks(self): # Register ourselves as a custom scanner check callbacks.registerScannerCheck(BurpScannerCheck()) # Register the proxy listener - montoya.proxy().registerRequestHandler(CustomProxyListener(app.scraped_headers)) + + # FIXME + # montoya.proxy().registerRequestHandler(CustomProxyListener(app.scraped_headers)) try: diff --git a/python/inql/scanner/customheaders.py b/python/inql/scanner/customheaders.py index fe29a3e..6453c7d 100644 --- a/python/inql/scanner/customheaders.py +++ b/python/inql/scanner/customheaders.py @@ -9,7 +9,7 @@ from ..globals import app from ..logger import log from ..utils.ui import inherits_popup_menu, ui_button, ui_label, ui_panel, ui_textarea - +from ..scraper.headers_scraper import HistoryScraper # from inql.actions.executor import ExecutorAction @@ -112,8 +112,12 @@ def __private_init__(self, text): # Create the set of custom headers associated to the session name app.custom_headers[text] = {} self._custom_headers = app.custom_headers[text] - self._scraped_headers = app.scraped_headers + + # FIXME Scraped headers does not need to be global anymore. + # self._scraped_headers = app.scraped_headers + self._scraped_headers = {} + # Data to store the state of the custom and scraped headers. # Inside the private data will be stored all the headers while in the # Data structures we will only store the selected ones (for the custom) @@ -132,6 +136,13 @@ def __private_init__(self, text): # Adding custom headers with object boolean self._augmenting_scraped_headers_data() + # Creting object to get the scraped headers (on demand) + self._header_scraper = HistoryScraper() + + # Data structure to hold removed/moved headers. + # For each domain it will hold the list of composed headers. + self._removed_scraped_headers = {} + return self def _build_domains_pane(self): @@ -280,25 +291,25 @@ def _domain_selection_listener(self): log.debug(new_header) self._custom_headers_dtm.addRow(new_header) - # get scraped domain + # get scraped headers self._scraped_headers_dtm.setRowCount(0) - if selected_domain in self._scraped_headers.keys(): - log.debug("Selected domain is in scraped headers") - for header in self._scraped_headers[selected_domain]: - log.debug("Scraped header to add is: %s: %s" % (header, self._scraped_headers[selected_domain][header])) - new_header = [] - new_header.append(header) - new_header.append(self._scraped_headers[selected_domain][header]) - # header = header.split(":") - # for elem in header: - # new_header.append(elem) - log.debug("New header to add is: ") - log.debug(new_header) - self._scraped_headers_dtm.addRow(new_header) + scraped_headers = self._header_scraper.get_scraped_headers(selected_domain) + if scraped_headers == None: + scraped_headers = [] - self._current_domain = selected_domain + if selected_domain not in self._removed_scraped_headers: + self._removed_scraped_headers[selected_domain] = set() + # removing all the "removed" headers from the list + for header in scraped_headers: + scraped_header = "{}: {}".format(header[0], header[1]) + log.debug("Header is: {}".format(scraped_header)) + if scraped_header not in self._removed_scraped_headers[selected_domain]: + self._scraped_headers_dtm.addRow(header) + + + self._current_domain = selected_domain def _add_custom_headers_row(self, _): @@ -331,17 +342,24 @@ def _remove_scraped_headers_row(self, _): """ rows = self._scraped_headers_table.getSelectedRows() for i in range(0, len(rows)): + name = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 0)) + value = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 1)) + scraped_header = "{}: {}".format(name, value) + + log.debug("The haders that has been removed is: ") + log.debug(scraped_header) + self._removed_scraped_headers[self._current_domain].add(scraped_header) self._scraped_headers_dtm.removeRow(rows[i] - i) - # TODO add scraped header modifier - nRow = self._scraped_headers_dtm.getRowCount() - log.debug("Removing all the scraped headers associated to this domain") - self._scraped_headers[self._current_domain] = {} + # # TODO add scraped header modifier + # nRow = self._scraped_headers_dtm.getRowCount() + # log.debug("Removing all the scraped headers associated to this domain") + # self._scraped_headers[self._current_domain] = {} - for i in range(0, nRow): - name = str(self._scraped_headers_dtm.getValueAt(i, 0)).lower() - value = str(self._scraped_headers_dtm.getValueAt(i, 1)).lower() - self._scraped_headers[self._current_domain][name] = value + # for i in range(0, nRow): + # name = str(self._scraped_headers_dtm.getValueAt(i, 0)).lower() + # value = str(self._scraped_headers_dtm.getValueAt(i, 1)).lower() + # self._scraped_headers[self._current_domain][name] = value @@ -354,15 +372,17 @@ def _move_scraped_headers_row(self, _): log.debug("The selected rows are:") log.debug(rows) - cols = self._scraped_headers_table.getColumnCount() - for i in range(0, len(rows)): - row_to_move = [False] - for j in range(cols): - row_to_move.append(self._scraped_headers_dtm.getValueAt(rows[i] - i, j)) + + name = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 0)) + value = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 1)) + scraped_header = "{}: {}".format(name, value) - log.debug("Adding new row: %s" % row_to_move) - self._custom_headers_dtm.addRow(row_to_move) + log.debug("The haders that has been moved is: ") + log.debug(scraped_header) + self._removed_scraped_headers[self._current_domain].add(scraped_header) + + self._custom_headers_dtm.addRow([name, value]) self._scraped_headers_dtm.removeRow(rows[i] - i) self._custom_headers_update() @@ -372,7 +392,7 @@ def _update_domains(self): domain table """ - for domain in self._scraped_headers: + for domain in self._header_scraper.get_scraped_domains(): log.debug("Considered domain: %s" % domain) if(domain != None and len(domain)>0): if domain in self._custom_headers: @@ -509,3 +529,5 @@ def _scraped_headers_update(self): # remove from data in case dest_data is false if self._dest_data == None: self._src_data.pop(new_row[1]) + + diff --git a/python/inql/scraper/headers_scraper.py b/python/inql/scraper/headers_scraper.py index 03dd460..29adb04 100644 --- a/python/inql/scraper/headers_scraper.py +++ b/python/inql/scraper/headers_scraper.py @@ -6,6 +6,7 @@ from ..globals import app from ..logger import log +from ..globals import montoya class CustomProxyListener(ProxyRequestHandler): @@ -66,3 +67,132 @@ def handleRequestReceived(self, interceptedRequest): def handleRequestToBeSent(self, interceptedRequest): return ProxyRequestToBeSentAction.continueWith(interceptedRequest.withDefaultHeaders()) + + +class HistoryScraper(): + + def __init__(self): + pass + + def get_scraped_headers(self, domain): + request_list = montoya.proxy().history() + + scraped_headers = [] + set_scraped_headers = set() + + for request in request_list: + + log.debug(request) + log.debug(type(request)) + + if domain != urlparse(request.finalRequest().url()).netloc: + # We want to filter request based on the domain + continue + + headers = request.finalRequest().headers() + + for header in headers: + + # removing connection header and host + # TODO Add more headers that we don't consider useful. + if header.name() == "Connection" or header.name() == "Host": + continue + + if header.value() == None or len(header.value()) <=0: + continue + + # Avoiding adding the same header (both name and value) twice. + scraped_header = "{}: {}".format(header.name(), header.value()) + if scraped_header not in set_scraped_headers: + set_scraped_headers.add(scraped_header) + scraped_headers.append([header.name().decode('utf-8'), header.value().decode('utf-8')]) + + return scraped_headers + + def get_scraped_domains(self): + request_list = montoya.proxy().history() + + domains = set() + + for request in request_list: + domains.add(urlparse(request.finalRequest().url()).netloc) + + + return sorted(domains) + + +# class CachedHistoryScraper(): + +# def __init__(self) -> None: +# # the scraped headers should track the "high level" session +# self.last_request_index = {} # session -> number +# self.scraped_headers = {} +# self.set_scraped_headers = {} + +# def get_scraped_headers(self, session) -> list[list[str, str]]: +# """ +# This function will return a list containing all the non identical headers seen +# in the history since the last query of this function. +# """ +# if session not in self.last_request_index: +# self.last_request_index[session] = 0 +# # self.scraped_headers[session] = [] +# # self.set_scraped_headers[session] = set() + +# request_list = montoya.proxy().history() +# request_list = request_list[self.last_request_index[session]:] + +# scraped_headers = [] +# set_scraped_headers = set() + +# for request in request_list: +# headers = request.finalRequest().headers() + +# for header in headers: + +# # removing connection header and host +# if header.name() == "Connection" or header.name() == "Host": +# continue + +# if header.value() == None or len(header.value()) <=0: +# continue + +# log.debug("The header is: ") +# log.debug(header.name()) +# log.debug(header.value()) +# log.debug() + +# scraped_header = "{}: {}".format(header.name(), header.value()) +# # if scraped_header not in self.set_scraped_headers[session]: +# # self.set_scraped_headers[session].add(scraped_header) +# # self.scraped_headers[session].append([header.name(), header.value()]) + +# if scraped_header not in set_scraped_headers: +# set_scraped_headers.add(scraped_header) +# scraped_headers.append([header.name().decode('utf-8'), header.value().decode('utf-8')]) + +# self.last_request_index[session] += len(request) +# return scraped_headers + + + + + +# TODO +""" +Scraped header should be based on domain and not on the session since we cannot control +(easily now) session switch and relate them with request. This means that scraped headers +will be agnostic from the session (which is use by the custom headers). + +However, scraped headers should be list[list[str,str]]. Indeed, a dictionary based on the header name +will easily lose all the headers with the same name but different parameters. + +To make this approach feasible, the most common headers should be removed such as: +- connection +- host +- path +- user-agent +- etc...7 + +""" + From 0e529a20ce92de6dae6e587d6a97cb880ec1d0d4 Mon Sep 17 00:00:00 2001 From: Matteo Oldani Date: Sat, 13 May 2023 22:41:28 +0200 Subject: [PATCH 2/3] Introduced sessions and fixed custom headers. --- python/inql/scanner/customheaders.py | 15 ++---- python/inql/scanner/introspection.py | 2 +- python/inql/scanner/omnibar.py | 71 +++++++++++++++++++++++++--- 3 files changed, 70 insertions(+), 18 deletions(-) diff --git a/python/inql/scanner/customheaders.py b/python/inql/scanner/customheaders.py index f00e86f..7f6a143 100644 --- a/python/inql/scanner/customheaders.py +++ b/python/inql/scanner/customheaders.py @@ -110,7 +110,9 @@ def __private_init__(self, text): self._current_domain = None # Create the set of custom headers associated to the session name - app.custom_headers[text] = {} + if text not in app.custom_headers: + app.custom_headers[text] = {} + self._custom_headers = app.custom_headers[text] @@ -254,7 +256,7 @@ def _augmenting_scraped_headers_data(self): self._scraped_headers_dtm.addRow(new_row) def _add_domain_listener(self, _): - name = JOptionPane.showInputDialog(self, "Enter domain name: ") + name = JOptionPane.showInputDialog(None, "Enter domain name: ", "New Domain", JOptionPane.INFORMATION_MESSAGE) if(name != None and len(name)>0): if name in self._custom_headers: log.info("You can't add the same domain twice") @@ -447,10 +449,6 @@ def _custom_headers_update(self): del self._custom_headers[self._current_domain][:] - # Create a global custom headers dictionary for the current domain - if self._current_domain not in app.custom_headers: - app.custom_headers[self._current_domain] = [] - nRow = self._custom_headers_dtm.getRowCount() nCol = self._custom_headers_dtm.getColumnCount() log.debug("Removing all the custom private data associated to this domain") @@ -477,11 +475,6 @@ def _custom_headers_update(self): self._custom_private_data[self._current_domain][idx] = new_row[0] log.debug("self._private_data[%s] = %s" % (new_row[1:], new_row[0])) - # Save the selected header in the global custom headers dictionary - app.custom_headers[self._current_domain].append(new_row[1:3]) - log.error("Saved the selected header in the global custom headers dictionary") - log.error("app.custom_headers[%s][%s] = %s" % (self._current_domain, new_row[1], new_row[2])) - # Adding the new row to the private headers to be displayed if new_row[0] == True: self._custom_headers[self._current_domain].append(new_row[1:]) diff --git a/python/inql/scanner/introspection.py b/python/inql/scanner/introspection.py index df981cb..7cc2a3c 100644 --- a/python/inql/scanner/introspection.py +++ b/python/inql/scanner/introspection.py @@ -116,7 +116,7 @@ def _analyze(url, filename=None, explicit_headers=None): # Create report directory (example: api.example.com/2023-02-15_102254) date = datetime.now().strftime("%Y-%m-%d_%H%M%S") - report_dir = "{}/{}".format(host, date) + report_dir = "{}/{}/{}".format(app.session_name, host, date) queries_dir = os.path.join(report_dir, 'queries') mutations_dir = os.path.join(report_dir, 'mutations') try: diff --git a/python/inql/scanner/omnibar.py b/python/inql/scanner/omnibar.py index 3576a4c..fa8b88b 100644 --- a/python/inql/scanner/omnibar.py +++ b/python/inql/scanner/omnibar.py @@ -7,7 +7,7 @@ from java.awt.event import ActionListener, FocusListener, KeyAdapter, KeyEvent from java.io import File from java.lang import System -from javax.swing import Box, BoxLayout, JFileChooser, JPanel, JSeparator, JTextField, SwingConstants, SwingUtilities +from javax.swing import Box, BoxLayout, JFileChooser, JOptionPane, JPanel, JSeparator, JTextField, SwingConstants, SwingUtilities, JComboBox, JButton, JMenuBar, JMenu, JMenuItem from ..editors.propertyeditor import SettingsEditor from ..globals import app @@ -18,6 +18,22 @@ from .customheaders import HeadersEditor from .introspection import analyze +class MenuActionListener(ActionListener): + """ + Custom Action Listener used to update the session selection. + """ + + def __init__(self, sessionmenu, item, name): + self.item = item + self.name = name + self.sessionmenu = sessionmenu + + def actionPerformed(self, event): + self.item.setSelected(True) + self.sessionmenu.setText(self.name + " " +u"\u25BC") + + # Change the global session identifier. + app.session_name = self.name class ScannerUrlField(FocusListener, KeyAdapter): """Textfield for the URL input. Shows a helpful hint when url is empty.""" @@ -29,16 +45,59 @@ def __init__(self, omnibar): log.debug("ScannerUrlField initiated") self.lock = Lock() self._omnibar = omnibar + self.session_names = set() super(ScannerUrlField, self).__init__() + def render(self): + + # Menu creation. + self.jmenu_bar= JMenuBar() + self.session_menu = JMenu("Sessions") + + self.default_menu_item = JMenuItem(app.session_name) + self.session_menu.add(self.default_menu_item, 0) + self.default_menu_item.addActionListener(MenuActionListener(self.session_menu, self.default_menu_item, app.session_name)) + + self.default_menu_item.setSelected(True) + self.session_menu.setText(app.session_name + " " + u'\u25BC') + + self.add_session_item = JMenuItem(" + Add Session") + self.add_session_item.addActionListener(self.addSessionActionListener) + + self.session_menu.add(self.add_session_item) + self.jmenu_bar.add(self.session_menu) + + # Omnibar creation. self.component = JTextField() self.component.setFocusable(True) self.component.putClientProperty("JTextField.placeholderText", self.hint) self.component.putClientProperty("JTextField.showClearButton", True) - self.component.addKeyListener(self) - return self.component + # Packing all together. + self.panel = JPanel(BorderLayout()) + self.panel.add(self.jmenu_bar, BorderLayout.WEST) + self.panel.add(self.component, BorderLayout.CENTER) + + return self.panel + + def addSessionActionListener(self, evnt): + """ + Action listener triggered when a new session needs to be created + """ + session_name = JOptionPane.showInputDialog(None, "Enter session Name", "New Session", JOptionPane.INFORMATION_MESSAGE) + if(session_name != None and len(session_name)>0): + if session_name in self.session_names: + log.info("You can't add the same domain twice") + return + self.session_names.add(session_name) + + new_session_item = JMenuItem(session_name) + self.session_menu.add(new_session_item, 0) + new_session_item.addActionListener(MenuActionListener(self.session_menu, new_session_item, session_name)) + + if session_name not in app.custom_headers: + app.custom_headers[session_name] = {} @single def keyPressed(self, e): @@ -299,9 +358,9 @@ def run(self): log.error("Current custom headers:") log.error(app.custom_headers) - if domain in app.custom_headers: + if domain in app.custom_headers[app.session_name]: log.debug("The URL has some custom headers set") - analyze(self.url, self.file, headers=app.custom_headers[domain]) + analyze(self.url, self.file, headers=app.custom_headers[app.session_name][domain]) else: log.debug("The URL has not set any custom headers, setting headers=none") analyze(self.url, self.file) @@ -360,7 +419,7 @@ def file(self, filename): self.file_field.value = filename def custom_header_button_handler(self, _): - HeadersEditor.get_instance() + HeadersEditor.get_instance(app.session_name) log.debug("Working") def settings_button_handler(self, _): From 049eb54a4ddc92f856e46942536a0d52252faa4b Mon Sep 17 00:00:00 2001 From: Matteo Oldani Date: Tue, 4 Jul 2023 18:32:49 +0200 Subject: [PATCH 3/3] Refactored custom headers --- python/inql/scanner/customheaders.py | 504 +++++++++++++++------------ 1 file changed, 284 insertions(+), 220 deletions(-) diff --git a/python/inql/scanner/customheaders.py b/python/inql/scanner/customheaders.py index 045349c..485b7ed 100644 --- a/python/inql/scanner/customheaders.py +++ b/python/inql/scanner/customheaders.py @@ -10,10 +10,6 @@ from ..scraper.headers_scraper import HistoryScraper from ..utils.pyswing import button -# from inql.actions.executor import ExecutorAction - - - class CustomTable(DefaultTableModel): """ @@ -82,7 +78,9 @@ def get_instance(text="Header Selector"): HeadersEditor.last_size = HeadersEditor.instances[text].this.getSize() # Before setting it as visible, update possible new domains - HeadersEditor.instances[text]._update_domains() + # HeadersEditor.instances[text]._update_domains() + HeadersEditor.instances[text].domain_pane.update() + # In any case I have to set it visible and on top HeadersEditor.instances[text].this.setVisible(True) @@ -109,16 +107,15 @@ def __private_init__(self, text): self._current_domain = None + # Creting object to get the scraped headers (on demand) + self._header_scraper = HistoryScraper() + # Create the set of custom headers associated to the session name if text not in app.custom_headers: app.custom_headers[text] = {} self._custom_headers = app.custom_headers[text] - - # FIXME Scraped headers does not need to be global anymore. - # self._scraped_headers = app.scraped_headers - self._scraped_headers = {} # Data to store the state of the custom and scraped headers. # Inside the private data will be stored all the headers while in the @@ -128,69 +125,17 @@ def __private_init__(self, text): # Table to display the Domains - self._build_domains_pane() + # self._build_domains_pane() DEPRECATED + + self.custom_header_tab = CustomHeaderTab(self, self._custom_headers) + self.scraped_header_tab = ScrapedHeadersTab(self, self._custom_headers) + self.domain_pane = DomainsPane(self, self._custom_headers, self._header_scraper, self.custom_header_tab, self.scraped_header_tab) - # Table to display the headers - self._build_custom_headers_pane() - self._build_scraped_headers_pane() + # # Table to display the headers self._build_gui_tabs() - # Adding custom headers with object boolean - self._augmenting_scraped_headers_data() - - # Creting object to get the scraped headers (on demand) - self._header_scraper = HistoryScraper() - - # Data structure to hold removed/moved headers. - # For each domain it will hold the list of composed headers. - self._removed_scraped_headers = {} - return self - def _build_domains_pane(self): - domain_colum = ["Domains"] - self._domain_table = JTable() - # self._domain_table.setDefaultEditor(object, None); - # self._domain_table.setDefaultRenderer(object, NonEditableColumnRenderer()) # Set the column renderer to non-editable - - self._domain_dtm = NonEditableModel(0,0) - self._domain_dtm.setColumnIdentifiers(domain_colum) - self._domain_table.setModel(self._domain_dtm) - self._domain_table.setSelectionMode(ListSelectionModel.SINGLE_SELECTION) - - self._domain_scroll_pane = JScrollPane(self._domain_table) - self._domain_scroll_pane.setPreferredSize(Dimension(150,200)) - - self._domain_table.getSelectionModel().addListSelectionListener(lambda _: self._domain_selection_listener()) - - self._add_domain_button = button("Add Domain", self._add_domain_listener, True) - self._domain_table_panel = JPanel(BorderLayout()) - self._domain_table_panel.add(self._domain_scroll_pane, BorderLayout.CENTER) - self._domain_table_panel.add(self._add_domain_button, BorderLayout.SOUTH) - - def _build_custom_headers_pane(self): - custom_headers_columns = ["Flag", "Header", "Value"] - self._custom_headers_table = JTable() - self._custom_headers_dtm = CustomTable(0, 0) - self._custom_headers_dtm.setColumnIdentifiers(custom_headers_columns) - self._custom_headers_table.setModel(self._custom_headers_dtm) - - self._custom_headers_dtm.addTableModelListener(lambda _: self._custom_headers_update()) - - # Create the "Add Row" button for the second table - self._add_custom_header = button("Add Header", self._add_custom_headers_row) - # Create the "Remove Row" button for the second table - self._remove_custom_header = button("Remove Headers", self._remove_custom_headers_row) - - # create the panel to hold the buttons - self._custom_headers_button_panel = JPanel(FlowLayout()) - self._custom_headers_button_panel.add(self._add_custom_header) - self._custom_headers_button_panel.add(self._remove_custom_header) - - self._custom_header_pane = JPanel(BorderLayout()) - self._custom_header_table_scroll_pane = JScrollPane(self._custom_headers_table) - self._custom_header_pane.add(self._custom_header_table_scroll_pane, BorderLayout.CENTER) - self._custom_header_pane.add(self._custom_headers_button_panel, BorderLayout.SOUTH); def _build_scraped_headers_pane(self): scraped_headers_columns = ["Header", "Value"] @@ -227,8 +172,8 @@ def _build_gui_tabs(self): self._scraped_headers_label = JLabel("Scraped Headers") self._scraped_headers_label.setHorizontalAlignment(JLabel.CENTER) - self._main_headers_panel.addTab("Custom Headers", self._custom_header_pane) - self._main_headers_panel.addTab("Scraped Headers", self._scraped_header_pane) + self._main_headers_panel.addTab("Custom Headers", self.custom_header_tab.get_custom_header_pane()) + self._main_headers_panel.addTab("Scraped Headers", self.scraped_header_tab.get_scraped_header_pane()) self._main_headers_panel.setTabLayoutPolicy(JTabbedPane.SCROLL_TAB_LAYOUT) self.this.addWindowListener(self) @@ -237,7 +182,7 @@ def _build_gui_tabs(self): self.this.setLayout(BorderLayout()) - self.this.add(self._domain_table_panel, BorderLayout.WEST) + self.this.add(self.domain_pane.get_domain_table_pane(), BorderLayout.WEST) self.this.add(self._main_headers_panel, BorderLayout.CENTER) self.this.pack() @@ -245,151 +190,142 @@ def _build_gui_tabs(self): self.this.setDefaultCloseOperation(JFrame.DO_NOTHING_ON_CLOSE) self.this.setVisible(True) - def _augmenting_custom_headers_data(self): - self._custom_headers_dtm.addRow(self._empty) - self._custom_headers_dtm.addRow(self._empty) - self._custom_headers_dtm.addRow(self._empty) + def windowClosing(self, _): + """ + Overrides WindowAdapter method - def _augmenting_scraped_headers_data(self): - for k in self._scraped_headers.keys(): - new_row = [k, self._scraped_headers[k]] - self._scraped_headers_dtm.addRow(new_row) + :param evt: unused + :return: None + """ - def _add_domain_listener(self, _): - name = JOptionPane.showInputDialog(None, "Enter domain name: ") - if(name != None and len(name)>0): - if name in self._custom_headers: - log.info("You can't add the same domain twice") - return + HeadersEditor.locations[self._text] = self.this.getLocation() + HeadersEditor.sizes[self._text] = self.this.getSize() + HeadersEditor.last_location = self.this.getLocation() + HeadersEditor.last_size = self.this.getSize() + HeadersEditor.offset = 0 + self.this.setVisible(False) - self._domain_dtm.addRow([name]) - self._custom_headers[name.encode('utf-8')] = [] # TODO check if it should be a dict or if a list is fine + # DEBUG + log.debug("Printing the headers at the moment of the custom headers window closing:") + log.debug(self._custom_headers) + for domain in self._custom_headers: + log.debug(self._custom_headers[domain]) + log.debug("Custom Headers in the app") + log.debug(app.custom_headers) - def _domain_selection_listener(self): - log.info("Domain selection listener") - # get selected domain - selected_row_number = self._domain_table.getSelectedRow() - selected_domain = str(self._domain_dtm.getValueAt(selected_row_number, 0)).lower() - log.debug("Selected row: " + selected_domain) - #update the current domain - self._current_domain = None + +class DomainsPane(): + """ + Class that defined the left panel which holds the domains. + """ - # get custom domains - self._custom_headers_dtm.setRowCount(0); - log.debug(self._custom_private_data.keys()) - if selected_domain in self._custom_private_data.keys(): - log.debug("The selected domain is in the custom private data") - # add the domains to the table - for header in self._custom_private_data[selected_domain]: - new_header = [] - new_header.append(self._custom_private_data[selected_domain][header]) - header = header.split(":") - for elem in header: - new_header.append(elem) + def __init__(self, global_obj, custom_headers, header_scraper, custom_header_tab, scraped_header_tab): - log.debug("New header to add is: ") - log.debug(new_header) - self._custom_headers_dtm.addRow(new_header) + domain_colum = ["Domains"] + self._domain_table = JTable() + + self._domain_dtm = NonEditableModel(0,0) + self._domain_dtm.setColumnIdentifiers(domain_colum) + self._domain_table.setModel(self._domain_dtm) + self._domain_table.setSelectionMode(ListSelectionModel.SINGLE_SELECTION) - # get scraped headers - self._scraped_headers_dtm.setRowCount(0) + self._domain_scroll_pane = JScrollPane(self._domain_table) + self._domain_scroll_pane.setPreferredSize(Dimension(150,200)) - scraped_headers = self._header_scraper.get_scraped_headers(selected_domain) - if scraped_headers == None: - scraped_headers = [] + self._domain_table.getSelectionModel().addListSelectionListener(lambda _: domain_selection_listener()) - if selected_domain not in self._removed_scraped_headers: - self._removed_scraped_headers[selected_domain] = set() - # removing all the "removed" headers from the list - for header in scraped_headers: - scraped_header = "{}: {}".format(header[0], header[1]) - log.debug("Header is: {}".format(scraped_header)) + self._custom_headers = custom_headers + self._header_scraper = header_scraper - if scraped_header not in self._removed_scraped_headers[selected_domain]: - self._scraped_headers_dtm.addRow(header) + self._custom_header_tab = custom_header_tab + self._scraped_header_tab = scraped_header_tab + self._global_obj = global_obj - self._current_domain = selected_domain + def add_domain_listener(_): + name = JOptionPane.showInputDialog(None, "Enter domain name: ") + if(name != None and len(name)>0): + if name in self._custom_headers: + log.info("You can't add the same domain twice") + return + self._domain_dtm.addRow([name]) + self._custom_headers[name.encode('utf-8')] = [] # TODO check if it should be a dict or if a list is fine - def _add_custom_headers_row(self, _): - """ - Add a new row the selection + def domain_selection_listener(): + log.info("Domain selection listener") - :return: None - """ - if(self._current_domain == None): - log.debug("You can't add a new line without having selected a domain") - return - self._custom_headers_dtm.addRow(self._empty) - self._custom_headers_update() + # get selected domain + selected_row_number = self._domain_table.getSelectedRow() + selected_domain = str(self._domain_dtm.getValueAt(selected_row_number, 0)).lower() + log.debug("Selected row: " + selected_domain) - def _remove_custom_headers_row(self, _): - """ - Remove all the selected rows from the selection - :return: - """ - rows = self._custom_headers_table.getSelectedRows() - for i in range(0, len(rows)): - self._custom_headers_dtm.removeRow(rows[i] - i) + #update the current domain + self._global_obj._current_domain = None - self._custom_headers_update() + # get custom domains + # self._custom_headers_dtm.setRowCount(0) + self._custom_header_tab.get_custom_header_table_model().setRowCount(0) - def _remove_scraped_headers_row(self, _): - """ - Remove all the selected rows from the selection - :return: - """ - rows = self._scraped_headers_table.getSelectedRows() - for i in range(0, len(rows)): - name = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 0)) - value = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 1)) - scraped_header = "{}: {}".format(name, value) + custom_private_data = self._custom_header_tab.get_custom_private_data() + log.debug(custom_private_data.keys()) + if selected_domain in custom_private_data.keys(): + log.debug("The selected domain is in the custom private data") + # add the domains to the table + for header in custom_private_data[selected_domain]: + new_header = [] + new_header.append(custom_private_data[selected_domain][header]) + header = header.split(":") + for elem in header: + new_header.append(elem) - log.debug("The haders that has been removed is: ") - log.debug(scraped_header) - self._removed_scraped_headers[self._current_domain].add(scraped_header) - self._scraped_headers_dtm.removeRow(rows[i] - i) + log.debug("New header to add is: ") + log.debug(new_header) + self._custom_header_tab.get_custom_header_table_model().addRow(new_header) - # # TODO add scraped header modifier - # nRow = self._scraped_headers_dtm.getRowCount() - # log.debug("Removing all the scraped headers associated to this domain") - # self._scraped_headers[self._current_domain] = {} + # get scraped headers + self._scraped_header_tab.get_scraped_header_table_model().setRowCount(0) + # self._scraped_headers_dtm.setRowCount(0) - # for i in range(0, nRow): - # name = str(self._scraped_headers_dtm.getValueAt(i, 0)).lower() - # value = str(self._scraped_headers_dtm.getValueAt(i, 1)).lower() - # self._scraped_headers[self._current_domain][name] = value + scraped_headers = self._header_scraper.get_scraped_headers(selected_domain) + if scraped_headers == None: + scraped_headers = [] + removed_scraped_header = self._scraped_header_tab.get_removed_scraped_headers() + if selected_domain not in removed_scraped_header: + removed_scraped_header[selected_domain] = set() + # removing all the "removed" headers from the list + for header in scraped_headers: + scraped_header = "{}: {}".format(header[0], header[1]) + log.debug("Header is: {}".format(scraped_header)) + if scraped_header not in removed_scraped_header[selected_domain]: + self._scraped_header_tab.get_scraped_header_table_model().addRow(header) - def _move_scraped_headers_row(self, _): - """ - Remove all the selected rows from the selection - :return: - """ - rows = self._scraped_headers_table.getSelectedRows() - log.debug("The selected rows are:") - log.debug(rows) - for i in range(0, len(rows)): - - name = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 0)) - value = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 1)) - scraped_header = "{}: {}".format(name, value) + self._global_obj._current_domain = selected_domain - log.debug("The haders that has been moved is: ") - log.debug(scraped_header) - self._removed_scraped_headers[self._current_domain].add(scraped_header) - self._custom_headers_dtm.addRow([name, value]) - self._scraped_headers_dtm.removeRow(rows[i] - i) - self._custom_headers_update() - def _update_domains(self): + self._add_domain_button = button("Add Domain", add_domain_listener, True) + self._domain_table_panel = JPanel(BorderLayout()) + self._domain_table_panel.add(self._domain_scroll_pane, BorderLayout.CENTER) + self._domain_table_panel.add(self._add_domain_button, BorderLayout.SOUTH) + + + def get_domain_table(self): + return self._domain_table + + def get_domain_table_model(self): + return self._domain_dtm + + def get_domain_table_pane(self): + return self._domain_table_panel + + def update(self): """ Checks the content of the Domains table and adds all the domains that are in the scraped headers but not in the domain table @@ -404,35 +340,82 @@ def _update_domains(self): self._domain_dtm.addRow([domain]) self._custom_headers[domain] = [] # TODO check if it should be a dict or if a list is fine + +class CustomHeaderTab(): + """ + Class that defines the custom header tab in the header editor tab. + """ - def windowClosing(self, _): - """ - Overrides WindowAdapter method + def __init__(self, global_obj, custom_headers): + custom_headers_columns = ["Flag", "Header", "Value"] + self._custom_headers_table = JTable() + self._custom_headers_dtm = CustomTable(0, 0) + self._custom_headers_dtm.setColumnIdentifiers(custom_headers_columns) + self._custom_headers_table.setModel(self._custom_headers_dtm) - :param evt: unused - :return: None - """ + self._custom_headers_dtm.addTableModelListener(lambda _: self.update()) - HeadersEditor.locations[self._text] = self.this.getLocation() - HeadersEditor.sizes[self._text] = self.this.getSize() - HeadersEditor.last_location = self.this.getLocation() - HeadersEditor.last_size = self.this.getSize() - HeadersEditor.offset = 0 - self.this.setVisible(False) + self._global_obj = global_obj - # DEBUG - log.debug("Printing the headers at the moment of the custom headers window closing:") - log.debug(self._custom_headers) - for domain in self._custom_headers: - log.debug(self._custom_headers[domain]) + self._custom_headers = custom_headers - log.debug("Custom Headers in the app") - log.debug(app.custom_headers) + # Data to store the state of the custom and scraped headers. + # Inside the private data will be stored all the headers while in the + # Data structures we will only store the selected ones (for the custom) + # For each domain we will store a dictionary + self._custom_private_data = {} + + self._empty = [False, "X-New-Header", "X-New-Header-Value"] + + def add_custom_headers_row(_): + + if(self._global_obj._current_domain == None): + log.debug("You can't add a new line without having selected a domain") + return + self._custom_headers_dtm.addRow(self._empty) + self.update() + + def remove_custom_headers_row(_): + rows = self._custom_headers_table.getSelectedRows() + for i in range(0, len(rows)): + self._custom_headers_dtm.removeRow(rows[i] - i) + self.update() + # Create the "Add Row" button for the second table + self._add_custom_header = button("Add Header", add_custom_headers_row, True) + # Create the "Remove Row" button for the second table + self._remove_custom_header = button("Remove Headers", remove_custom_headers_row) + + # create the panel to hold the buttons + self._custom_headers_button_panel = JPanel(FlowLayout()) + self._custom_headers_button_panel.add(self._add_custom_header) + self._custom_headers_button_panel.add(self._remove_custom_header) + + self._custom_header_pane = JPanel(BorderLayout()) + self._custom_header_table_scroll_pane = JScrollPane(self._custom_headers_table) + self._custom_header_pane.add(self._custom_header_table_scroll_pane, BorderLayout.CENTER) + self._custom_header_pane.add(self._custom_headers_button_panel, BorderLayout.SOUTH); + + + def get_custom_header_table(self): + return self._custom_headers_table + - def _custom_headers_update(self): + def get_custom_header_table_model(self): + return self._custom_headers_dtm + + + def get_custom_header_pane(self): + return self._custom_header_pane + + + def get_custom_private_data(self): + return self._custom_private_data + + + def update(self): """ Update the data content with the updated rows @@ -442,17 +425,17 @@ def _custom_headers_update(self): :return: None """ - if(self._current_domain == None): + if(self._global_obj._current_domain == None): log.debug("You can't add a new line without having selected a domain") return - del self._custom_headers[self._current_domain][:] + del self._custom_headers[self._global_obj._current_domain][:] nRow = self._custom_headers_dtm.getRowCount() nCol = self._custom_headers_dtm.getColumnCount() log.debug("Removing all the custom private data associated to this domain") - self._custom_private_data[self._current_domain] = {} + self._custom_private_data[self._global_obj._current_domain] = {} for i in range(0, nRow): # Creating the new row new_row = [None] * nCol @@ -472,17 +455,92 @@ def _custom_headers_update(self): idx = "%s:%s" % (new_row[1], new_row[2]) log.debug("The idx is: %s" % idx) - self._custom_private_data[self._current_domain][idx] = new_row[0] + self._custom_private_data[self._global_obj._current_domain][idx] = new_row[0] log.debug("self._private_data[%s] = %s" % (new_row[1:], new_row[0])) # Save the selected header in the global custom headers dictionary - app.custom_headers[self._current_domain].append(new_row[1:3]) + self._custom_headers[self._global_obj._current_domain].append(new_row[1:3]) # TODO this was app.custom_header, check if correct log.debug("Saved the selected header in the global custom headers dictionary") - log.debug("app.custom_headers[%s][%s] = %s" % (self._current_domain, new_row[1], new_row[2])) + log.debug("app.custom_headers[%s][%s] = %s" % (self._global_obj._current_domain, new_row[1], new_row[2])) # Adding the new row to the private headers to be displayed if new_row[0] == True: - self._custom_headers[self._current_domain].append(new_row[1:]) + self._custom_headers[self._global_obj._current_domain].append(new_row[1:]) + +class ScrapedHeadersTab(): + + def __init__(self, global_obj, custom_header): + scraped_headers_columns = ["Header", "Value"] + self._scraped_headers_table = JTable() + self._scraped_headers_dtm = CustomTable(0, 0) + self._scraped_headers_dtm.setColumnIdentifiers(scraped_headers_columns) + self._scraped_headers_table.setModel(self._scraped_headers_dtm) + + self._scraped_headers = {} + + # Data structure to hold removed/moved headers. + # For each domain it will hold the list of composed headers. + self._removed_scraped_headers = {} + + self._global_obj = global_obj + + + def remove_scraped_headers_row(_): + """ + Remove all the selected rows from the selection + :return: + """ + rows = self._scraped_headers_table.getSelectedRows() + for i in range(0, len(rows)): + name = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 0)) + value = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 1)) + scraped_header = "{}: {}".format(name, value) + + log.debug("The haders that has been removed is: ") + log.debug(scraped_header) + self._removed_scraped_headers[self._global_obj._current_domain].add(scraped_header) + self._scraped_headers_dtm.removeRow(rows[i] - i) + + def move_scraped_headers_row(_): + """ + Remove all the selected rows from the selection + :return: + """ + rows = self._scraped_headers_table.getSelectedRows() + log.debug("The selected rows are:") + log.debug(rows) + + for i in range(0, len(rows)): + + name = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 0)) + value = str(self._scraped_headers_dtm.getValueAt(rows[i] - i, 1)) + scraped_header = "{}: {}".format(name, value) + + log.debug("The haders that has been moved is: ") + log.debug(scraped_header) + self._removed_scraped_headers[self._global_obj._current_domain].add(scraped_header) + + custom_header.get_custom_header_table_model().addRow([name, value]) + # self._custom_headers_dtm.addRow([name, value]) + self._scraped_headers_dtm.removeRow(rows[i] - i) + self._custom_headers_update() + + self._move_scraped_headers = button("Move Headers", move_scraped_headers_row) + self._remove_scraped_headers = button("Remove Headers", remove_scraped_headers_row) + + self._scraped_headers_button_panel = JPanel(FlowLayout()) + self._scraped_headers_button_panel.add(self._move_scraped_headers) + self._scraped_headers_button_panel.add(self._remove_scraped_headers) + + self._scraped_header_pane = JPanel(BorderLayout()) + self._scraped_header_table_scroll_pane = JScrollPane(self._scraped_headers_table) + self._scraped_header_pane.add(self._scraped_header_table_scroll_pane, BorderLayout.CENTER) + self._scraped_header_pane.add(self._scraped_headers_button_panel, BorderLayout.SOUTH) + + for k in self._scraped_headers.keys(): + new_row = [k, self._scraped_headers[k]] + self._scraped_headers_dtm.addRow(new_row) + def _scraped_headers_update(self): """ @@ -494,13 +552,8 @@ def _scraped_headers_update(self): :return: None """ - # log.debug("Updating the model") - # if self._dest_data: - # del self._dest_data[:] - # else: - # del self._src_data[:] - if(self._current_domain == None): + if(self._global_obj._current_domain == None): log.debug("You can't add a new line without having selected a domain") return @@ -525,8 +578,6 @@ def _scraped_headers_update(self): idx = "%s:%s" % (new_row[1], new_row[2]) log.debug("The idx is: %s" % idx) - self._private_data[idx] = new_row[0] - log.debug("self._private_data[%s] = %s" % (new_row[1:], new_row[0])) # Adding the new row to the private headers to be displayed if new_row[0] == True: @@ -539,4 +590,17 @@ def _scraped_headers_update(self): if self._dest_data == None: self._src_data.pop(new_row[1]) - + + def get_scraped_header_table(self): + return self._scraped_headers_table + + + def get_scraped_header_table_model(self): + return self._scraped_headers_dtm + + + def get_scraped_header_pane(self): + return self._scraped_header_pane + + def get_removed_scraped_headers(self): + return self._removed_scraped_headers \ No newline at end of file