def test_index():
screen = pyte.HistoryScreen(5, 5, history=50)
# Filling the screen with line numbers, so it's easier to
# track history contents.
for idx in range(screen.lines):
screen.draw(str(idx))
if idx != screen.lines - 1:
screen.linefeed()
assert not screen.history.top
assert not screen.history.bottom
# a) first index, expecting top history to be updated.
line = screen.buffer[0]
screen.index()
assert screen.history.top
assert screen.history.top[-1] == line
# b) second index.
line = screen.buffer[0]
screen.index()
assert len(screen.history.top) == 2
assert screen.history.top[-1] == line
# c) rotation.
for _ in range(screen.history.size * 2):
screen.index()
assert len(screen.history.top) == 50
def test_ensure_width(monkeypatch):
screen = pyte.HistoryScreen(5, 5, history=50)
screen.set_mode(mo.LNM)
escape = dict(pyte.Stream.escape)
escape.update({"N": "next_page", "P": "prev_page"})
monkeypatch.setattr(pyte.Stream, "escape", escape)
stream = pyte.Stream(screen)
for idx in range(len(screen.buffer) * 5):
stream.feed(str(idx) + os.linesep)
assert screen.display == ["21 ", "22 ", "23 ", "24 ", " "]
# a) shrinking the screen, expecting the lines displayed to
# be truncated.
screen.resize(5, 2)
stream.feed(ctrl.ESC + "P")
assert all(len(l) != 2 for l in screen.history.top)
assert all(len(l) == 2 for l in screen.history.bottom)
assert screen.display == ["18", "19", "20", "21", "22"]
# b) expading the screen, expecting the lines displayed to
# be filled with whitespace characters.
screen.resize(5, 10)
stream.feed(ctrl.ESC + "N")
assert all(len(l) == 10 for l in list(screen.history.top)[-3:])
assert all(len(l) != 10 for l in screen.history.bottom)
assert screen.display == [
"21 ", "22 ", "23 ", "24 ", " "
]
def test_reverse_index():
screen = pyte.HistoryScreen(5, 5, history=50)
# Filling the screen with line numbers, so it's easier to
# track history contents.
for idx in range(len(screen.buffer)):
screen.draw(str(idx))
if idx != len(screen.buffer) - 1:
screen.linefeed()
assert not screen.history.top
assert not screen.history.bottom
screen.cursor_position()
# a) first index, expecting top history to be updated.
line = screen.buffer[-1]
screen.reverse_index()
assert screen.history.bottom
assert screen.history.bottom[0] == line
# b) second index.
line = screen.buffer[-1]
screen.reverse_index()
assert len(screen.history.bottom) == 2
assert screen.history.bottom[1] == line
# c) rotation.
for _ in range(len(screen.buffer)**screen.lines):
screen.reverse_index()
assert len(screen.history.bottom) == 50
def __init__(self, columns, lines, p_in):
self.screen = pyte.HistoryScreen(columns, lines)
self.screen.set_mode(pyte.modes.LNM)
self.screen.write_process_input = \
lambda data: p_in.write(data.encode())
self.stream = pyte.ByteStream()
self.stream.attach(self.screen)
def test_not_enough_lines():
screen = pyte.HistoryScreen(5, 5, history=6)
screen.set_mode(mo.LNM)
for idx in range(screen.lines):
screen.draw(str(idx))
screen.linefeed()
assert screen.history.top
assert not screen.history.bottom
assert screen.history.position == 6
assert screen.display == ["1 ", "2 ", "3 ", "4 ", " "]
screen.prev_page()
assert not screen.history.top
assert len(screen.history.bottom) is 1
assert chars(screen.history.bottom, screen.columns) == [" "]
assert screen.display == [
"0 ",
"1 ",
"2 ",
"3 ",
"4 ",
]
screen.next_page()
assert screen.history.top
assert not screen.history.bottom
assert screen.display == ["1 ", "2 ", "3 ", "4 ", " "]
def test_ensure_width(monkeypatch):
screen = pyte.HistoryScreen(5, 5, history=50)
screen.set_mode(mo.LNM)
escape = dict(pyte.Stream.escape)
escape.update({"N": "next_page", "P": "prev_page"})
monkeypatch.setattr(pyte.Stream, "escape", escape)
stream = pyte.Stream(screen)
for idx in range(screen.lines * 5):
stream.feed("{0:04d}".format(idx) + os.linesep)
assert screen.display == ["0021 ", "0022 ", "0023 ", "0024 ", " "]
# Shrinking the screen should truncate the displayed lines following lines.
screen.resize(5, 3)
stream.feed(ctrl.ESC + "P")
# Inequality because we have an all-empty last line.
assert all(len(l) <= 3 for l in screen.history.bottom)
assert screen.display == [
"001", # 18
"001", # 19
"002", # 20
"002", # 21
"002" # 22
]
def process_ssh(data):
msg_type = data[1]
sid = get_session_id(data)
msg = data[18:]
session = sessions[sid]
if msg_type == 0x00: # connection
ip = '.'.join(str(i) for i in msg[:4])
port = msg[4] * 256 + msg[5]
session['data'] = {
'sessionID': sid,
'svrType': 'ssh',
'potIP': POT_IP,
'portPort': 22,
'attackIP': ip,
'attackPort': port,
'blastList': [],
'execCommands': [],
'startTime': datetime.datetime.now().isoformat(),
'resize': 0,
'strokes': []
}
elif msg_type == 0x01: # login
username_len = msg[0]
username = str(msg[1:username_len + 1])
password_len = msg[username_len + 1]
password = str(msg[username_len + 2:])
session['data']['blastList'].append('%s::%s' % (username, password))
elif msg_type == 0x02: # session
pair = session['data']['blastList'][-1].split('::')
session['data']['loginName'] = pair[0]
session['data']['loginPasswd'] = pair[1]
screen = pyte.HistoryScreen(80, 24, history=1000, ratio=1)
stream = pyte.ByteStream(screen)
session['screen'] = screen
session['stream'] = stream
elif msg_type == 0x03: # window resize
r = msg[0]
c = msg[1]
session['screen'].resize(r, c)
session['data']['resize'] += 1
elif msg_type == 0x04: # data message
session['stream'].feed(msg)
session['data']['strokes'].append(msg)
elif msg_type == 0x05: # disconnection
session['data']['endTime'] = datetime.datetime.now().isoformat()
if 'screen' in session:
screen = session['screen']
while screen.history.position > screen.lines and screen.history.top:
screen.prev_page()
data = ''
rest_lines = 0
while screen.history.position < screen.history.size and screen.history.bottom:
data += '\n'.join(screen.display) + '\n'
rest_lines = len(screen.history.bottom)
screen.next_page()
data += '\n'.join(screen.display[-rest_lines:]) + '\n'
sha256 = send_content(data)
send(session['data'], sha256, 'ssh')
del sessions[sid]
def __init__(self, window_size=None, newline="\r\n",
encoding=('UTF-8', 'ignore')):
self._rows, self._columns = window_size or (200, 200)
self._newline = newline
self._stream = pyte.ByteStream(encodings=[encoding])
self._screen = pyte.HistoryScreen(self._rows,
self._columns,
history=100000)
self._stream.attach(self._screen)
self._screen.set_charset('B', '(')
self._buffer = ''
self._whitespace_after_last_feed = ''
def test_erase_in_display():
screen = pyte.HistoryScreen(5, 5, history=6)
screen.set_mode(mo.LNM)
for idx in range(screen.lines):
screen.draw(str(idx))
screen.linefeed()
screen.prev_page()
# See #80 on GitHub for details.
screen.erase_in_display(3)
assert not screen.history.top
assert not screen.history.bottom
def test_next_page():
screen = pyte.HistoryScreen(5, 5, history=50)
screen.set_mode(mo.LNM)
# Once again filling the screen with line numbers, but this time,
# we need them to span on multiple lines.
for idx in range(len(screen.buffer) * 5):
screen.draw(str(idx))
screen.linefeed()
assert screen.history.top
assert not screen.history.bottom
assert screen.history.position == 50
assert len(screen.buffer) == screen.lines
assert screen.display == ["21 ", "22 ", "23 ", "24 ", " "]
# a) page up -- page down.
screen.prev_page()
screen.next_page()
assert screen.history.top
assert not screen.history.bottom
assert screen.history.position == 50
assert screen.display == ["21 ", "22 ", "23 ", "24 ", " "]
# b) double page up -- page down.
screen.prev_page()
screen.prev_page()
screen.next_page()
assert screen.history.position == 45
assert screen.history.top
assert chars(screen.history.bottom) == ["23 ", "24 ", " "]
assert len(screen.buffer) == screen.lines
assert screen.display == ["18 ", "19 ", "20 ", "21 ", "22 "]
# c) double page up -- double page down
screen.prev_page()
screen.prev_page()
screen.next_page()
screen.next_page()
assert screen.history.position == 45
assert len(screen.buffer) == screen.lines
assert screen.display == ["18 ", "19 ", "20 ", "21 ", "22 "]
def start_screen_mode(self):
""" Starts the ``screen mode``.
In the ``screen mode``, the output is just the same with the real terminal. It
means that any real-time application likes ``top`` will be captured as-is.
Consecutive `read` from this VChannel instance may produce redundancy ouput.
"""
channel = self._channels[self._current_name]
channel['screen'] = pyte.HistoryScreen(channel['w'], channel['h'],
100000)
channel['stream'] = pyte.Stream()
channel['stream'].attach(channel['screen'])
# handle different version of pyte
try:
channel['screen'].set_charset('B', '(')
except:
channel['screen'].define_charset('B', '(')
BuiltIn().log("Started ``screen mode``")
def start_screen_mode(self):
""" Starts the ``screen mode``.
In the ``screen mode``, the output is just the same with the real terminal. It
means that any real-time application likes ``top`` will be captured as-is.
Consecutive `read` from this VChannel instance may produce redundancy ouput.
"""
channel = self._channels[self._current_name]
if not channel['screen']:
channel['screen'] = pyte.HistoryScreen(channel['w'], channel['h'],
100000)
# channel['screen'] = pyte.Screen(channel['w'], channel['h'])
# channel['stream'] = pyte.ByteStream(encodings=[('UTF-8', 'ignore')])
channel['stream'] = pyte.Stream()
channel['stream'].attach(channel['screen'])
channel['screen'].set_charset('B', '(')
BuiltIn().log("Started ``screen mode``")
def test_cursor_is_hidden(monkeypatch):
screen = pyte.HistoryScreen(5, 5, history=50)
escape = dict(pyte.Stream.escape)
escape.update({"N": "next_page", "P": "prev_page"})
monkeypatch.setattr(pyte.Stream, "escape", escape)
stream = pyte.Stream(screen)
for idx in range(screen.lines * 5):
stream.feed(str(idx) + os.linesep)
assert not screen.cursor.hidden
stream.feed(ctrl.ESC + "P")
assert screen.cursor.hidden
stream.feed(ctrl.ESC + "P")
assert screen.cursor.hidden
stream.feed(ctrl.ESC + "N")
assert screen.cursor.hidden
stream.feed(ctrl.ESC + "N")
assert not screen.cursor.hidden
def test_draw(monkeypatch):
screen = pyte.HistoryScreen(5, 5, history=50)
screen.set_mode(mo.LNM)
escape = dict(pyte.Stream.escape)
escape.update({"N": "next_page", "P": "prev_page"})
monkeypatch.setattr(pyte.Stream, "escape", escape)
stream = pyte.Stream(screen)
for idx in range(screen.lines * 5):
stream.feed(str(idx) + os.linesep)
assert screen.display == ["21 ", "22 ", "23 ", "24 ", " "]
# a) doing a pageup and then a draw -- expecting the screen
# to scroll to the bottom before drawing anything.
stream.feed(ctrl.ESC + "P")
stream.feed(ctrl.ESC + "P")
stream.feed(ctrl.ESC + "N")
stream.feed("x")
assert screen.display == ["21 ", "22 ", "23 ", "24 ", "x "]
def __init__(self) -> None:
self.screen = pyte.HistoryScreen(60, 20, history=9999, ratio = 0.05)
self.screen.set_mode(pyte.modes.LNM)
self.stream = pyte.ByteStream(self.screen)
def test_prev_page():
screen = pyte.HistoryScreen(4, 4, history=40)
screen.set_mode(mo.LNM)
assert screen.history.position == 40
# Once again filling the screen with line numbers, but this time,
# we need them to span on multiple lines.
for idx in range(screen.lines * 10):
screen.draw(str(idx))
screen.linefeed()
assert screen.history.top
assert not screen.history.bottom
assert screen.history.position == 40
assert screen.display == ["37 ", "38 ", "39 ", " "]
assert chars(screen.history.top,
screen.columns)[-4:] == ["33 ", "34 ", "35 ", "36 "]
# a) first page up.
screen.prev_page()
assert screen.history.position == 38
assert len(screen.buffer) == screen.lines
assert screen.display == ["35 ", "36 ", "37 ", "38 "]
assert chars(screen.history.top,
screen.columns)[-4:] == ["31 ", "32 ", "33 ", "34 "]
assert len(screen.history.bottom) == 2
assert chars(screen.history.bottom, screen.columns) == [
"39 ",
" ",
]
# b) second page up.
screen.prev_page()
assert screen.history.position == 36
assert len(screen.buffer) == screen.lines
assert screen.display == [
"33 ",
"34 ",
"35 ",
"36 ",
]
assert len(screen.history.bottom) == 4
assert chars(screen.history.bottom, screen.columns) == [
"37 ",
"38 ",
"39 ",
" ",
]
# c) same with odd number of lines.
screen = pyte.HistoryScreen(5, 5, history=50)
screen.set_mode(mo.LNM)
for idx in range(screen.lines * 10):
screen.draw(str(idx))
screen.linefeed()
assert screen.history.top
assert not screen.history.bottom
assert screen.history.position == 50
assert screen.display == ["46 ", "47 ", "48 ", "49 ", " "]
screen.prev_page()
assert screen.history.position == 47
assert screen.display == ["43 ", "44 ", "45 ", "46 ", "47 "]
assert len(screen.history.bottom) == 3
assert chars(screen.history.bottom, screen.columns) == [
"48 ",
"49 ",
" ",
]
# d) with a ratio other than 0.5
screen = pyte.HistoryScreen(4, 4, history=40, ratio=0.75)
screen.set_mode(mo.LNM)
for idx in range(screen.lines * 10):
screen.draw(str(idx))
screen.linefeed()
assert screen.history.top
assert not screen.history.bottom
assert screen.history.position == 40
assert screen.display == ["37 ", "38 ", "39 ", " "]
screen.prev_page()
assert screen.history.position == 37
assert screen.display == ["34 ", "35 ", "36 ", "37 "]
assert len(screen.history.bottom) == 3
assert chars(screen.history.bottom,
screen.columns) == ["38 ", "39 ", " "]
# e) same with cursor in the middle of the screen.
screen = pyte.HistoryScreen(5, 5, history=50)
screen.set_mode(mo.LNM)
for idx in range(screen.lines * 10):
screen.draw(str(idx))
screen.linefeed()
assert screen.history.top
assert not screen.history.bottom
assert screen.history.position == 50
assert screen.display == ["46 ", "47 ", "48 ", "49 ", " "]
screen.cursor_to_line(screen.lines // 2)
while screen.history.position > screen.lines:
screen.prev_page()
assert screen.history.position == screen.lines
assert len(screen.buffer) == screen.lines
assert screen.display == ["1 ", "2 ", "3 ", "4 ", "5 "]
while screen.history.position < screen.history.size:
screen.next_page()
assert screen.history.position == screen.history.size
assert screen.display == ["46 ", "47 ", "48 ", "49 ", " "]
# e) same with cursor near the middle of the screen.
screen = pyte.HistoryScreen(5, 5, history=50)
screen.set_mode(mo.LNM)
for idx in range(screen.lines * 10):
screen.draw(str(idx))
screen.linefeed()
assert screen.history.top
assert not screen.history.bottom
assert screen.history.position == 50
assert screen.display == ["46 ", "47 ", "48 ", "49 ", " "]
screen.cursor_to_line(screen.lines // 2 - 2)
while screen.history.position > screen.lines:
screen.prev_page()
assert screen.history.position == screen.lines
assert len(screen.buffer) == screen.lines
assert screen.display == ["1 ", "2 ", "3 ", "4 ", "5 "]
while screen.history.position < screen.history.size:
screen.next_page()
assert screen.history.position == screen.history.size
assert len(screen.buffer) == screen.lines
assert screen.display == ["46 ", "47 ", "48 ", "49 ", " "]
print(pyte.control.ESC + pyte.escape.RIS)
for idx, line in enumerate(screen.display, 1):
print("{0:2d} {1} ¶".format(idx, line))
input(os.linesep + os.linesep + text)
def random_string(n, alphabet=string.ascii_letters + " "):
return "".join(random.choice(alphabet) for _ in range(n))
if __name__ == "__main__":
stream = pyte.Stream()
# ``ratio=1`` means scroll the whole screen.
screen = pyte.HistoryScreen(80, 12, ratio=1)
screen.set_mode(mo.LNM)
stream.attach(screen)
pages = 3
stream.feed(
os.linesep.join(
random_string(screen.columns)
for _ in range(screen.lines * pages)))
screen.prev_page()
print_screen(screen, "Hit ENTER to move up!")
screen.prev_page()
print_screen(screen, "Hit ENTER to move back down!")
screen.next_page()
print_screen(screen, "OK?")
