Skip to content

Commit

Permalink
FIXED: New prompt_toolkit syntax
Browse files Browse the repository at this point in the history
  • Loading branch information
@shramos
shramos committed May 21, 2018
1 parent 88102e4 commit 791e620
Show file tree
Hide file tree
Showing 5 changed files with 86 additions and 66 deletions.
21 changes: 11 additions & 10 deletions polymorph/UI/fieldinterface.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,12 @@
# For more information about the project: https://github.com/shramos/polymorph

from polymorph.UI.interface import Interface
from prompt_toolkit import prompt
from prompt_toolkit import PromptSession
from prompt_toolkit import HTML
from collections import OrderedDict
from polymorph.UI.command_parser import CommandParser
from prompt_toolkit.history import FileHistory
from prompt_toolkit.contrib.completers import WordCompleter
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.shortcuts import CompleteStyle

Expand Down Expand Up @@ -45,15 +45,15 @@ def run(self):
completer = WordCompleter(['value', 'type', 'show', 'name', 'slice',
'custom', 'size', 'dump', 'clear', 'back'])
history = FileHistory(self._polym_path + '/.finterface_history')
session = PromptSession(history=history)
while True:
try:
command = prompt(HTML("<bold>PH:cap/t%d/%s/<red>%s</red> > </bold>" %
(self._tindex, self._lname, self._f.name)),
history=history,
completer=completer,
complete_style=CompleteStyle.READLINE_LIKE,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True)
command = session.prompt(HTML("<bold>PH:cap/t%d/%s/<red>%s</red> > </bold>" %
(self._tindex, self._lname, self._f.name)),
completer=completer,
complete_style=CompleteStyle.READLINE_LIKE,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True)
except KeyboardInterrupt:
self.exit_program()
continue
Expand Down Expand Up @@ -117,7 +117,8 @@ def _type(self, command):
self._f.to_int(args["-o"])
Interface._print_info("New type added")
else:
Interface._print_error("Wrong order. Please select big or little")
Interface._print_error(
"Wrong order. Please select big or little")
else:
Interface._print_error(
"Wrong type. Please choose between ('hex', 'bytes', 'str', 'int')")
Expand Down
39 changes: 23 additions & 16 deletions polymorph/UI/layerinterface.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
# For more information about the project: https://github.com/shramos/polymorph

from polymorph.UI.interface import Interface
from prompt_toolkit import prompt
from prompt_toolkit import PromptSession
from prompt_toolkit import HTML
from collections import OrderedDict
from polymorph.UI.command_parser import CommandParser
Expand All @@ -12,7 +12,7 @@
from polymorph.UI.fieldinterface import FieldInterface
import hexdump
from prompt_toolkit.history import FileHistory
from prompt_toolkit.contrib.completers import WordCompleter
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.shortcuts import CompleteStyle
import construct
Expand Down Expand Up @@ -47,15 +47,15 @@ def run(self):
completer = WordCompleter(['show', 'name', 'field', 'fields',
'dump', 'recalculate', 'clear', 'back'])
history = FileHistory(self._polym_path + '/.linterface_history')
session = PromptSession(history=history)
while True:
try:
command = prompt(HTML("<bold>PH:cap/t%d/<red>%s</red> > </bold>" %
(self._tindex, self._l.name)),
history=history,
completer=completer,
complete_style=CompleteStyle.READLINE_LIKE,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True)
command = session.prompt(HTML("<bold>PH:cap/t%d/<red>%s</red> > </bold>" %
(self._tindex, self._l.name)),
completer=completer,
complete_style=CompleteStyle.READLINE_LIKE,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True)
except KeyboardInterrupt:
self.exit_program()
continue
Expand Down Expand Up @@ -217,7 +217,8 @@ def _field(self, command):
if len(command) == 1:
Interface.print_help(LayerInterface._field_help())
elif len(command) == 2 and command[1].lower() in self._l.fieldnames():
fi = FieldInterface(self._l.getfield(command[1].lower()), self._tindex, self._l.name, self._poisoner)
fi = FieldInterface(self._l.getfield(
command[1].lower()), self._tindex, self._l.name, self._poisoner)
fi.run()
else:
cp = CommandParser(LayerInterface._field_opts())
Expand All @@ -234,7 +235,8 @@ def _field(self, command):
start = input("Start byte of the custom field: ")
end = input("End byte of the custom field: ")
if not start.isdecimal() or not end.isdecimal():
Interface._print_error("The start or end byte is not a number")
Interface._print_error(
"The start or end byte is not a number")
return
else:
fslice = slice(int(start), int(end))
Expand All @@ -259,15 +261,18 @@ def _field(self, command):
new_field.to_bytes()
# Add the field to the layer
self._l.addfield(new_field)
Interface._print_info("Field %s added to the layer" % args['-a'])
Interface._print_info(
"Field %s added to the layer" % args['-a'])
# Deletes a field from the layer
elif args["-d"]:
del_field = self._l.getfield(args["-d"])
if del_field:
self._l.delfield(del_field)
Interface._print_info("Field %s deleted from the layer" % args["-d"])
Interface._print_info(
"Field %s deleted from the layer" % args["-d"])
else:
Interface._print_error("The field %s is not in the layer" % args["-d"])
Interface._print_error(
"The field %s is not in the layer" % args["-d"])

@staticmethod
def _field_help():
Expand Down Expand Up @@ -366,8 +371,10 @@ def _recalculate(self, command):
fields = LayerInterface._extrac_deps(args["-sb"], args["-e"])
if fields:
try:
self._l.add_struct(args["-f"], fields, args["-sb"], args["-e"])
Interface._print_info("Struct added to field %s" % args["-f"])
self._l.add_struct(
args["-f"], fields, args["-sb"], args["-e"])
Interface._print_info(
"Struct added to field %s" % args["-f"])
except:
Interface._print_error(
"Wrong fields or wrong syntax referring to the fields")
Expand Down
19 changes: 9 additions & 10 deletions polymorph/UI/maininterface.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
# For more information about the project: https://github.com/shramos/polymorph

from polymorph.UI.interface import Interface
from prompt_toolkit import prompt
from prompt_toolkit import PromptSession
from prompt_toolkit.shortcuts import CompleteStyle
from prompt_toolkit import HTML
from polymorph.utils import capture, get_arpspoofer, set_ip_forwarding, readtemplate, readpcap
Expand All @@ -12,7 +12,7 @@
from collections import OrderedDict
from polymorph.UI.command_parser import CommandParser
from prompt_toolkit.history import FileHistory
from prompt_toolkit.contrib.completers import WordCompleter
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
import os

Expand All @@ -31,14 +31,14 @@ def run(self):
"""Runs the interface and waits for user input commands."""
completer = WordCompleter(['capture', 'spoof', 'clear', 'import'])
history = FileHistory(self._polym_path + '/.minterface_history')
session = PromptSession(history=history)
while True:
try:
command = prompt(HTML("<bold><red>PH</red> > </bold>"),
history=history,
completer=completer,
complete_style=CompleteStyle.READLINE_LIKE,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True)
command = session.prompt(HTML("<bold><red>PH</red> > </bold>"),
completer=completer,
complete_style=CompleteStyle.READLINE_LIKE,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True)
except KeyboardInterrupt:
self.exit_program()
continue
Expand Down Expand Up @@ -216,12 +216,11 @@ def _import_help():
("-h", "prints the help."),
("-t", "path to a template to be imported."),
("-pcap", "path to a pcap file to be imported")

])
return OrderedDict([
("name", "import"),
("usage", "import [-option]"),
("description", "Import different objects in the framework, such as templates or captures."),
("options", options)
])

54 changes: 33 additions & 21 deletions polymorph/UI/templateinterface.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,10 @@
# For more information about the project: https://github.com/shramos/polymorph

from polymorph.UI.interface import Interface
from prompt_toolkit import prompt, HTML
from prompt_toolkit import PromptSession, HTML
from prompt_toolkit.shortcuts import confirm, CompleteStyle
from prompt_toolkit.history import FileHistory
from prompt_toolkit.contrib.completers import WordCompleter
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from collections import OrderedDict
from polymorph.UI.command_parser import CommandParser
Expand Down Expand Up @@ -56,15 +56,15 @@ def run(self):
'description', 'spoof', 'clear', 'back'])
# Initialization of the command history
history = FileHistory(join(self._polym_path, '.tinterface_history'))
session = PromptSession(history=history)
while True:
try:
command = prompt(HTML("<bold>PH:cap/<red>t%d</red> > </bold>" %
self._index),
history=history,
completer=completer,
complete_style=CompleteStyle.READLINE_LIKE,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True)
command = session.prompt(HTML("<bold>PH:cap/<red>t%d</red> > </bold>" %
self._index),
completer=completer,
complete_style=CompleteStyle.READLINE_LIKE,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True)
except KeyboardInterrupt:
self.exit_program()
continue
Expand Down Expand Up @@ -242,7 +242,8 @@ def _conditions(self, command, cond):
else:
Interface._print_error("Wrong index")
except ValueError:
Interface._print_error("Please enter a positive integer")
Interface._print_error(
"Please enter a positive integer")
# Deletes a condition
if args['-d']:
try:
Expand All @@ -257,9 +258,11 @@ def _conditions(self, command, cond):
# Create the new file if not exists
if not os.path.isfile(join(self._conds_path, cond, args["-a"] + ".py")):
self._create_cond(cond, args["-a"])
ret = os.system("%s %s.py" % (args["-e"], join(self._conds_path, cond, args["-a"])))
ret = os.system("%s %s.py" % (
args["-e"], join(self._conds_path, cond, args["-a"])))
if ret != 0:
Interface._print_error("The editor is not installed or is not in the PATH")
Interface._print_error(
"The editor is not installed or is not in the PATH")
return
# Save the condition to the Template
try:
Expand All @@ -269,9 +272,11 @@ def _conditions(self, command, cond):
"Bad syntax, please check the code syntax")
return
# If user wants, delete condition from disk
keepindisk = confirm('Keep file on disk (%s)? [y/N] ' % join(self._conds_path, cond))
keepindisk = confirm(
'Keep file on disk (%s)? [y/N] ' % join(self._conds_path, cond))
if not keepindisk:
os.remove("%s.py" % join(self._conds_path, cond, args["-a"]))
os.remove("%s.py" %
join(self._conds_path, cond, args["-a"]))
Interface._print_info("Condition %s added" % args["-a"])
# Imports a condition from disk
elif args['-i']:
Expand All @@ -287,7 +292,8 @@ def _conditions(self, command, cond):
self._add_cond(cond, name)
Interface._print_info("Condition %s imported" % args['-i'])
except ModuleNotFoundError:
Interface._print_error("The condition %s is not in disk" % args['-i'])
Interface._print_error(
"The condition %s is not in disk" % args['-i'])
print("(Please place your .py file in correct path)")
return

Expand All @@ -299,7 +305,8 @@ def _create_cond(self, cond, name):

def _add_cond(self, cond, name):
"""Adds a new condition to the `Template`."""
m = importlib.import_module("polymorph.conditions.%s.%s" % (cond, name))
m = importlib.import_module(
"polymorph.conditions.%s.%s" % (cond, name))
importlib.reload(m)
self._t.add_function(cond, name, getattr(m, dir(m)[-1]))

Expand Down Expand Up @@ -518,7 +525,8 @@ def _save(self, command):
self._t.write(args['-p'])
Interface._print_info("Template saved to disk")
except:
Interface._print_error("The path %s does not exist" % args['-p'])
Interface._print_error(
"The path %s does not exist" % args['-p'])

@staticmethod
def _save_help():
Expand Down Expand Up @@ -549,7 +557,8 @@ def _layer(self, command):
if len(command) == 1:
Interface.print_help(TemplateInterface._layer_help())
elif len(command) == 2 and command[1].upper() in self._t.layernames():
li = LayerInterface(self._t.getlayer(command[1].upper()), self._index, self._poisoner)
li = LayerInterface(self._t.getlayer(
command[1].upper()), self._index, self._poisoner)
li.run()
else:
cp = CommandParser(TemplateInterface._layer_opts())
Expand All @@ -571,17 +580,20 @@ def _layer(self, command):
new_layer = TLayer(args["-a"], raw=self._t.raw.hex(),
lslice=lslice, custom=True)
self._t.addlayer(new_layer)
Interface._print_info("New layer %s added to the Template" % args["-a"])
Interface._print_info(
"New layer %s added to the Template" % args["-a"])
else:
Interface._print_error("The start or end byte is not a number")
Interface._print_error(
"The start or end byte is not a number")
# Deletes an existing layer
elif args["-d"]:
del_layer = self._t.getlayer(args["-d"])
if del_layer:
self._t.dellayer(del_layer)
Interface._print_info("Layer %s deleted" % args["-d"])
else:
Interface._print_error("The layer %s does not exist" % args["-d"])
Interface._print_error(
"The layer %s does not exist" % args["-d"])

@staticmethod
def _layer_help():
Expand Down
19 changes: 10 additions & 9 deletions polymorph/UI/tlistinterface.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,14 @@
# For more information about the project: https://github.com/shramos/polymorph

from polymorph.UI.interface import Interface
from prompt_toolkit import prompt
from prompt_toolkit import PromptSession
from prompt_toolkit import HTML
from collections import OrderedDict
from polymorph.UI.command_parser import CommandParser
from polymorph.UI.templateinterface import TemplateInterface
import os
from prompt_toolkit.history import FileHistory
from prompt_toolkit.contrib.completers import WordCompleter
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.shortcuts import CompleteStyle
import polymorph
Expand Down Expand Up @@ -44,14 +44,14 @@ def run(self):
completer = WordCompleter(['show', 'dissect', 'template', 'wireshark',
'clear', 'back'])
history = FileHistory(self._polym_path + '/.tlinterface_history')
session = PromptSession(history=history)
while True:
try:
command = prompt(HTML("<bold>PH:<red>cap</red> > </bold>"),
history=history,
completer=completer,
complete_style=CompleteStyle.READLINE_LIKE,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True)
command = session.prompt(HTML("<bold>PH:<red>cap</red> > </bold>"),
completer=completer,
complete_style=CompleteStyle.READLINE_LIKE,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True)
except KeyboardInterrupt:
self.exit_program()
continue
Expand Down Expand Up @@ -198,7 +198,8 @@ def _wireshark(self, command):
"""Opens Wireshark with the actual `Template` List in pcap format."""
if len(command) == 1:
Interface._print_info("Opening Wireshark...")
os.system("nohup wireshark %s &" % join(self._polym_path, ".tmp.pcap"))
os.system("nohup wireshark %s &" %
join(self._polym_path, ".tmp.pcap"))
return
# Parsing arguments
cp = CommandParser(TListInterface._wireshark_opts())
Expand Down

0 comments on commit 791e620

Please sign in to comment.