564 lines
22 KiB
Python
564 lines
22 KiB
Python
|
# Copyright 2014 Google Inc. All rights reserved.
|
||
|
|
#
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
#
|
||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
#
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
"""ADB protocol implementation.
|
||
|
|
|
||
|
|
Implements the ADB protocol as seen in android's adb/adbd binaries, but only the
|
||
|
|
host side.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import struct
|
||
|
|
import time
|
||
|
|
from io import BytesIO
|
||
|
|
from adb import usb_exceptions
|
||
|
|
|
||
|
|
# Maximum amount of data in an ADB packet.
|
||
|
|
MAX_ADB_DATA = 4096
|
||
|
|
# ADB protocol version.
|
||
|
|
VERSION = 0x01000000
|
||
|
|
|
||
|
|
# AUTH constants for arg0.
|
||
|
|
AUTH_TOKEN = 1
|
||
|
|
AUTH_SIGNATURE = 2
|
||
|
|
AUTH_RSAPUBLICKEY = 3
|
||
|
|
|
||
|
|
|
||
|
|
def find_backspace_runs(stdout_bytes, start_pos):
|
||
|
|
first_backspace_pos = stdout_bytes[start_pos:].find(b'\x08')
|
||
|
|
if first_backspace_pos == -1:
|
||
|
|
return -1, 0
|
||
|
|
|
||
|
|
end_backspace_pos = (start_pos + first_backspace_pos) + 1
|
||
|
|
while True:
|
||
|
|
if chr(stdout_bytes[end_backspace_pos]) == '\b':
|
||
|
|
end_backspace_pos += 1
|
||
|
|
else:
|
||
|
|
break
|
||
|
|
|
||
|
|
num_backspaces = end_backspace_pos - (start_pos + first_backspace_pos)
|
||
|
|
|
||
|
|
return (start_pos + first_backspace_pos), num_backspaces
|
||
|
|
|
||
|
|
|
||
|
|
class InvalidCommandError(Exception):
|
||
|
|
"""Got an invalid command over USB."""
|
||
|
|
|
||
|
|
def __init__(self, message, response_header, response_data):
|
||
|
|
if response_header == b'FAIL':
|
||
|
|
message = 'Command failed, device said so. (%s)' % message
|
||
|
|
super(InvalidCommandError, self).__init__(
|
||
|
|
message, response_header, response_data)
|
||
|
|
|
||
|
|
|
||
|
|
class InvalidResponseError(Exception):
|
||
|
|
"""Got an invalid response to our command."""
|
||
|
|
|
||
|
|
|
||
|
|
class InvalidChecksumError(Exception):
|
||
|
|
"""Checksum of data didn't match expected checksum."""
|
||
|
|
|
||
|
|
|
||
|
|
class InterleavedDataError(Exception):
|
||
|
|
"""We only support command sent serially."""
|
||
|
|
|
||
|
|
|
||
|
|
def MakeWireIDs(ids):
|
||
|
|
id_to_wire = {
|
||
|
|
cmd_id: sum(c << (i * 8) for i, c in enumerate(bytearray(cmd_id)))
|
||
|
|
for cmd_id in ids
|
||
|
|
}
|
||
|
|
wire_to_id = {wire: cmd_id for cmd_id, wire in id_to_wire.items()}
|
||
|
|
return id_to_wire, wire_to_id
|
||
|
|
|
||
|
|
|
||
|
|
class AuthSigner(object):
|
||
|
|
"""Signer for use with authenticated ADB, introduced in 4.4.x/KitKat."""
|
||
|
|
|
||
|
|
def Sign(self, data):
|
||
|
|
"""Signs given data using a private key."""
|
||
|
|
raise NotImplementedError()
|
||
|
|
|
||
|
|
def GetPublicKey(self):
|
||
|
|
"""Returns the public key in PEM format without headers or newlines."""
|
||
|
|
raise NotImplementedError()
|
||
|
|
|
||
|
|
|
||
|
|
class _AdbConnection(object):
|
||
|
|
"""ADB Connection."""
|
||
|
|
|
||
|
|
def __init__(self, usb, local_id, remote_id, timeout_ms):
|
||
|
|
self.usb = usb
|
||
|
|
self.local_id = local_id
|
||
|
|
self.remote_id = remote_id
|
||
|
|
self.timeout_ms = timeout_ms
|
||
|
|
|
||
|
|
def _Send(self, command, arg0, arg1, data=b''):
|
||
|
|
message = AdbMessage(command, arg0, arg1, data)
|
||
|
|
message.Send(self.usb, self.timeout_ms)
|
||
|
|
|
||
|
|
def Write(self, data):
|
||
|
|
"""Write a packet and expect an Ack."""
|
||
|
|
self._Send(b'WRTE', arg0=self.local_id, arg1=self.remote_id, data=data)
|
||
|
|
# Expect an ack in response.
|
||
|
|
cmd, okay_data = self.ReadUntil(b'OKAY')
|
||
|
|
if cmd != b'OKAY':
|
||
|
|
if cmd == b'FAIL':
|
||
|
|
raise usb_exceptions.AdbCommandFailureException(
|
||
|
|
'Command failed.', okay_data)
|
||
|
|
raise InvalidCommandError(
|
||
|
|
'Expected an OKAY in response to a WRITE, got %s (%s)',
|
||
|
|
cmd, okay_data)
|
||
|
|
return len(data)
|
||
|
|
|
||
|
|
def Okay(self):
|
||
|
|
self._Send(b'OKAY', arg0=self.local_id, arg1=self.remote_id)
|
||
|
|
|
||
|
|
def ReadUntil(self, *expected_cmds):
|
||
|
|
"""Read a packet, Ack any write packets."""
|
||
|
|
cmd, remote_id, local_id, data = AdbMessage.Read(
|
||
|
|
self.usb, expected_cmds, self.timeout_ms)
|
||
|
|
if local_id != 0 and self.local_id != local_id:
|
||
|
|
raise InterleavedDataError("We don't support multiple streams...")
|
||
|
|
if remote_id != 0 and self.remote_id != remote_id:
|
||
|
|
raise InvalidResponseError(
|
||
|
|
'Incorrect remote id, expected %s got %s' % (
|
||
|
|
self.remote_id, remote_id))
|
||
|
|
# Ack write packets.
|
||
|
|
if cmd == b'WRTE':
|
||
|
|
self.Okay()
|
||
|
|
return cmd, data
|
||
|
|
|
||
|
|
def ReadUntilClose(self):
|
||
|
|
"""Yield packets until a Close packet is received."""
|
||
|
|
while True:
|
||
|
|
cmd, data = self.ReadUntil(b'CLSE', b'WRTE')
|
||
|
|
if cmd == b'CLSE':
|
||
|
|
self._Send(b'CLSE', arg0=self.local_id, arg1=self.remote_id)
|
||
|
|
break
|
||
|
|
if cmd != b'WRTE':
|
||
|
|
if cmd == b'FAIL':
|
||
|
|
raise usb_exceptions.AdbCommandFailureException(
|
||
|
|
'Command failed.', data)
|
||
|
|
raise InvalidCommandError('Expected a WRITE or a CLOSE, got %s (%s)',
|
||
|
|
cmd, data)
|
||
|
|
yield data
|
||
|
|
|
||
|
|
def Close(self):
|
||
|
|
self._Send(b'CLSE', arg0=self.local_id, arg1=self.remote_id)
|
||
|
|
cmd, data = self.ReadUntil(b'CLSE')
|
||
|
|
if cmd != b'CLSE':
|
||
|
|
if cmd == b'FAIL':
|
||
|
|
raise usb_exceptions.AdbCommandFailureException('Command failed.', data)
|
||
|
|
raise InvalidCommandError('Expected a CLSE response, got %s (%s)',
|
||
|
|
cmd, data)
|
||
|
|
|
||
|
|
|
||
|
|
class AdbMessage(object):
|
||
|
|
"""ADB Protocol and message class.
|
||
|
|
|
||
|
|
Protocol Notes
|
||
|
|
|
||
|
|
local_id/remote_id:
|
||
|
|
Turns out the documentation is host/device ambidextrous, so local_id is the
|
||
|
|
id for 'the sender' and remote_id is for 'the recipient'. So since we're
|
||
|
|
only on the host, we'll re-document with host_id and device_id:
|
||
|
|
|
||
|
|
OPEN(host_id, 0, 'shell:XXX')
|
||
|
|
READY/OKAY(device_id, host_id, '')
|
||
|
|
WRITE(0, host_id, 'data')
|
||
|
|
CLOSE(device_id, host_id, '')
|
||
|
|
"""
|
||
|
|
|
||
|
|
ids = [b'SYNC', b'CNXN', b'AUTH', b'OPEN', b'OKAY', b'CLSE', b'WRTE']
|
||
|
|
commands, constants = MakeWireIDs(ids)
|
||
|
|
# An ADB message is 6 words in little-endian.
|
||
|
|
format = b'<6I'
|
||
|
|
|
||
|
|
connections = 0
|
||
|
|
|
||
|
|
def __init__(self, command=None, arg0=None, arg1=None, data=b''):
|
||
|
|
self.command = self.commands[command]
|
||
|
|
self.magic = self.command ^ 0xFFFFFFFF
|
||
|
|
self.arg0 = arg0
|
||
|
|
self.arg1 = arg1
|
||
|
|
self.data = data
|
||
|
|
|
||
|
|
@property
|
||
|
|
def checksum(self):
|
||
|
|
return self.CalculateChecksum(self.data)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def CalculateChecksum(data):
|
||
|
|
# The checksum is just a sum of all the bytes. I swear.
|
||
|
|
if isinstance(data, bytearray):
|
||
|
|
total = sum(data)
|
||
|
|
elif isinstance(data, bytes):
|
||
|
|
if data and isinstance(data[0], bytes):
|
||
|
|
# Python 2 bytes (str) index as single-character strings.
|
||
|
|
total = sum(map(ord, data))
|
||
|
|
else:
|
||
|
|
# Python 3 bytes index as numbers (and PY2 empty strings sum() to 0)
|
||
|
|
total = sum(data)
|
||
|
|
else:
|
||
|
|
# Unicode strings (should never see?)
|
||
|
|
total = sum(map(ord, data))
|
||
|
|
return total & 0xFFFFFFFF
|
||
|
|
|
||
|
|
def Pack(self):
|
||
|
|
"""Returns this message in an over-the-wire format."""
|
||
|
|
return struct.pack(self.format, self.command, self.arg0, self.arg1,
|
||
|
|
len(self.data), self.checksum, self.magic)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def Unpack(cls, message):
|
||
|
|
try:
|
||
|
|
cmd, arg0, arg1, data_length, data_checksum, unused_magic = struct.unpack(
|
||
|
|
cls.format, message)
|
||
|
|
except struct.error as e:
|
||
|
|
raise ValueError('Unable to unpack ADB command.', cls.format, message, e)
|
||
|
|
return cmd, arg0, arg1, data_length, data_checksum
|
||
|
|
|
||
|
|
def Send(self, usb, timeout_ms=None):
|
||
|
|
"""Send this message over USB."""
|
||
|
|
usb.BulkWrite(self.Pack(), timeout_ms)
|
||
|
|
usb.BulkWrite(self.data, timeout_ms)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def Read(cls, usb, expected_cmds, timeout_ms=None, total_timeout_ms=None):
|
||
|
|
"""Receive a response from the device."""
|
||
|
|
total_timeout_ms = usb.Timeout(total_timeout_ms)
|
||
|
|
start = time.time()
|
||
|
|
while True:
|
||
|
|
msg = usb.BulkRead(24, timeout_ms)
|
||
|
|
cmd, arg0, arg1, data_length, data_checksum = cls.Unpack(msg)
|
||
|
|
command = cls.constants.get(cmd)
|
||
|
|
if not command:
|
||
|
|
raise InvalidCommandError(
|
||
|
|
'Unknown command: %x' % cmd, cmd, (arg0, arg1))
|
||
|
|
if command in expected_cmds:
|
||
|
|
break
|
||
|
|
|
||
|
|
if time.time() - start > total_timeout_ms:
|
||
|
|
raise InvalidCommandError(
|
||
|
|
'Never got one of the expected responses (%s)' % expected_cmds,
|
||
|
|
cmd, (timeout_ms, total_timeout_ms))
|
||
|
|
|
||
|
|
if data_length > 0:
|
||
|
|
data = bytearray()
|
||
|
|
while data_length > 0:
|
||
|
|
temp = usb.BulkRead(data_length, timeout_ms)
|
||
|
|
if len(temp) != data_length:
|
||
|
|
print(
|
||
|
|
"Data_length {} does not match actual number of bytes read: {}".format(data_length, len(temp)))
|
||
|
|
data += temp
|
||
|
|
|
||
|
|
data_length -= len(temp)
|
||
|
|
|
||
|
|
actual_checksum = cls.CalculateChecksum(data)
|
||
|
|
if actual_checksum != data_checksum:
|
||
|
|
raise InvalidChecksumError(
|
||
|
|
'Received checksum %s != %s', (actual_checksum, data_checksum))
|
||
|
|
else:
|
||
|
|
data = b''
|
||
|
|
return command, arg0, arg1, bytes(data)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def Connect(cls, usb, banner=b'notadb', rsa_keys=None, auth_timeout_ms=100):
|
||
|
|
"""Establish a new connection to the device.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
usb: A USBHandle with BulkRead and BulkWrite methods.
|
||
|
|
banner: A string to send as a host identifier.
|
||
|
|
rsa_keys: List of AuthSigner subclass instances to be used for
|
||
|
|
authentication. The device can either accept one of these via the Sign
|
||
|
|
method, or we will send the result of GetPublicKey from the first one
|
||
|
|
if the device doesn't accept any of them.
|
||
|
|
auth_timeout_ms: Timeout to wait for when sending a new public key. This
|
||
|
|
is only relevant when we send a new public key. The device shows a
|
||
|
|
dialog and this timeout is how long to wait for that dialog. If used
|
||
|
|
in automation, this should be low to catch such a case as a failure
|
||
|
|
quickly; while in interactive settings it should be high to allow
|
||
|
|
users to accept the dialog. We default to automation here, so it's low
|
||
|
|
by default.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The device's reported banner. Always starts with the state (device,
|
||
|
|
recovery, or sideload), sometimes includes information after a : with
|
||
|
|
various product information.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
usb_exceptions.DeviceAuthError: When the device expects authentication,
|
||
|
|
but we weren't given any valid keys.
|
||
|
|
InvalidResponseError: When the device does authentication in an
|
||
|
|
unexpected way.
|
||
|
|
"""
|
||
|
|
# In py3, convert unicode to bytes. In py2, convert str to bytes.
|
||
|
|
# It's later joined into a byte string, so in py2, this ends up kind of being a no-op.
|
||
|
|
if isinstance(banner, str):
|
||
|
|
banner = bytearray(banner, 'utf-8')
|
||
|
|
|
||
|
|
msg = cls(
|
||
|
|
command=b'CNXN', arg0=VERSION, arg1=MAX_ADB_DATA,
|
||
|
|
data=b'host::%s\0' % banner)
|
||
|
|
msg.Send(usb)
|
||
|
|
cmd, arg0, arg1, banner = cls.Read(usb, [b'CNXN', b'AUTH'])
|
||
|
|
if cmd == b'AUTH':
|
||
|
|
if not rsa_keys:
|
||
|
|
raise usb_exceptions.DeviceAuthError(
|
||
|
|
'Device authentication required, no keys available.')
|
||
|
|
# Loop through our keys, signing the last 'banner' or token.
|
||
|
|
for rsa_key in rsa_keys:
|
||
|
|
if arg0 != AUTH_TOKEN:
|
||
|
|
raise InvalidResponseError(
|
||
|
|
'Unknown AUTH response: %s %s %s' % (arg0, arg1, banner))
|
||
|
|
|
||
|
|
# Do not mangle the banner property here by converting it to a string
|
||
|
|
signed_token = rsa_key.Sign(banner)
|
||
|
|
msg = cls(
|
||
|
|
command=b'AUTH', arg0=AUTH_SIGNATURE, arg1=0, data=signed_token)
|
||
|
|
msg.Send(usb)
|
||
|
|
cmd, arg0, unused_arg1, banner = cls.Read(usb, [b'CNXN', b'AUTH'])
|
||
|
|
if cmd == b'CNXN':
|
||
|
|
return banner
|
||
|
|
# None of the keys worked, so send a public key.
|
||
|
|
msg = cls(
|
||
|
|
command=b'AUTH', arg0=AUTH_RSAPUBLICKEY, arg1=0,
|
||
|
|
data=rsa_keys[0].GetPublicKey() + b'\0')
|
||
|
|
msg.Send(usb)
|
||
|
|
try:
|
||
|
|
cmd, arg0, unused_arg1, banner = cls.Read(
|
||
|
|
usb, [b'CNXN'], timeout_ms=auth_timeout_ms)
|
||
|
|
except usb_exceptions.ReadFailedError as e:
|
||
|
|
if e.usb_error.value == -7: # Timeout.
|
||
|
|
raise usb_exceptions.DeviceAuthError(
|
||
|
|
'Accept auth key on device, then retry.')
|
||
|
|
raise
|
||
|
|
# This didn't time-out, so we got a CNXN response.
|
||
|
|
return banner
|
||
|
|
return banner
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def Open(cls, usb, destination, timeout_ms=None):
|
||
|
|
"""Opens a new connection to the device via an OPEN message.
|
||
|
|
|
||
|
|
Not the same as the posix 'open' or any other google3 Open methods.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
usb: USB device handle with BulkRead and BulkWrite methods.
|
||
|
|
destination: The service:command string.
|
||
|
|
timeout_ms: Timeout in milliseconds for USB packets.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
InvalidResponseError: Wrong local_id sent to us.
|
||
|
|
InvalidCommandError: Didn't get a ready response.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The local connection id.
|
||
|
|
"""
|
||
|
|
local_id = 1
|
||
|
|
msg = cls(
|
||
|
|
command=b'OPEN', arg0=local_id, arg1=0,
|
||
|
|
data=destination + b'\0')
|
||
|
|
msg.Send(usb, timeout_ms)
|
||
|
|
cmd, remote_id, their_local_id, _ = cls.Read(usb, [b'CLSE', b'OKAY'],
|
||
|
|
timeout_ms=timeout_ms)
|
||
|
|
if local_id != their_local_id:
|
||
|
|
raise InvalidResponseError(
|
||
|
|
'Expected the local_id to be {}, got {}'.format(local_id, their_local_id))
|
||
|
|
if cmd == b'CLSE':
|
||
|
|
# Some devices seem to be sending CLSE once more after a request, this *should* handle it
|
||
|
|
cmd, remote_id, their_local_id, _ = cls.Read(usb, [b'CLSE', b'OKAY'],
|
||
|
|
timeout_ms=timeout_ms)
|
||
|
|
# Device doesn't support this service.
|
||
|
|
if cmd == b'CLSE':
|
||
|
|
return None
|
||
|
|
if cmd != b'OKAY':
|
||
|
|
raise InvalidCommandError('Expected a ready response, got {}'.format(cmd),
|
||
|
|
cmd, (remote_id, their_local_id))
|
||
|
|
return _AdbConnection(usb, local_id, remote_id, timeout_ms)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def Command(cls, usb, service, command='', timeout_ms=None):
|
||
|
|
"""One complete set of USB packets for a single command.
|
||
|
|
|
||
|
|
Sends service:command in a new connection, reading the data for the
|
||
|
|
response. All the data is held in memory, large responses will be slow and
|
||
|
|
can fill up memory.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
usb: USB device handle with BulkRead and BulkWrite methods.
|
||
|
|
service: The service on the device to talk to.
|
||
|
|
command: The command to send to the service.
|
||
|
|
timeout_ms: Timeout for USB packets, in milliseconds.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
InterleavedDataError: Multiple streams running over usb.
|
||
|
|
InvalidCommandError: Got an unexpected response command.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The response from the service.
|
||
|
|
"""
|
||
|
|
return ''.join(cls.StreamingCommand(usb, service, command, timeout_ms))
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def StreamingCommand(cls, usb, service, command='', timeout_ms=None):
|
||
|
|
"""One complete set of USB packets for a single command.
|
||
|
|
|
||
|
|
Sends service:command in a new connection, reading the data for the
|
||
|
|
response. All the data is held in memory, large responses will be slow and
|
||
|
|
can fill up memory.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
usb: USB device handle with BulkRead and BulkWrite methods.
|
||
|
|
service: The service on the device to talk to.
|
||
|
|
command: The command to send to the service.
|
||
|
|
timeout_ms: Timeout for USB packets, in milliseconds.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
InterleavedDataError: Multiple streams running over usb.
|
||
|
|
InvalidCommandError: Got an unexpected response command.
|
||
|
|
|
||
|
|
Yields:
|
||
|
|
The responses from the service.
|
||
|
|
"""
|
||
|
|
if not isinstance(command, bytes):
|
||
|
|
command = command.encode('utf8')
|
||
|
|
connection = cls.Open(
|
||
|
|
usb, destination=b'%s:%s' % (service, command),
|
||
|
|
timeout_ms=timeout_ms)
|
||
|
|
for data in connection.ReadUntilClose():
|
||
|
|
yield data.decode('utf8')
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def InteractiveShellCommand(cls, conn, cmd=None, strip_cmd=True, delim=None, strip_delim=True, clean_stdout=True):
|
||
|
|
"""Retrieves stdout of the current InteractiveShell and sends a shell command if provided
|
||
|
|
TODO: Should we turn this into a yield based function so we can stream all output?
|
||
|
|
|
||
|
|
Args:
|
||
|
|
conn: Instance of AdbConnection
|
||
|
|
cmd: Optional. Command to run on the target.
|
||
|
|
strip_cmd: Optional (default True). Strip command name from stdout.
|
||
|
|
delim: Optional. Delimiter to look for in the output to know when to stop expecting more output
|
||
|
|
(usually the shell prompt)
|
||
|
|
strip_delim: Optional (default True): Strip the provided delimiter from the output
|
||
|
|
clean_stdout: Cleanup the stdout stream of any backspaces and the characters that were deleted by the backspace
|
||
|
|
Returns:
|
||
|
|
The stdout from the shell command.
|
||
|
|
"""
|
||
|
|
|
||
|
|
if delim is not None and not isinstance(delim, bytes):
|
||
|
|
delim = delim.encode('utf-8')
|
||
|
|
|
||
|
|
# Delimiter may be shell@hammerhead:/ $
|
||
|
|
# The user or directory could change, making the delimiter somthing like root@hammerhead:/data/local/tmp $
|
||
|
|
# Handle a partial delimiter to search on and clean up
|
||
|
|
if delim:
|
||
|
|
user_pos = delim.find(b'@')
|
||
|
|
dir_pos = delim.rfind(b':/')
|
||
|
|
if user_pos != -1 and dir_pos != -1:
|
||
|
|
partial_delim = delim[user_pos:dir_pos + 1] # e.g. @hammerhead:
|
||
|
|
else:
|
||
|
|
partial_delim = delim
|
||
|
|
else:
|
||
|
|
partial_delim = None
|
||
|
|
|
||
|
|
stdout = ''
|
||
|
|
stdout_stream = BytesIO()
|
||
|
|
original_cmd = ''
|
||
|
|
|
||
|
|
try:
|
||
|
|
|
||
|
|
if cmd:
|
||
|
|
original_cmd = str(cmd)
|
||
|
|
cmd += '\r' # Required. Send a carriage return right after the cmd
|
||
|
|
cmd = cmd.encode('utf8')
|
||
|
|
|
||
|
|
# Send the cmd raw
|
||
|
|
bytes_written = conn.Write(cmd)
|
||
|
|
|
||
|
|
if delim:
|
||
|
|
# Expect multiple WRTE cmds until the delim (usually terminal prompt) is detected
|
||
|
|
|
||
|
|
data = b''
|
||
|
|
while partial_delim not in data:
|
||
|
|
cmd, data = conn.ReadUntil(b'WRTE')
|
||
|
|
stdout_stream.write(data)
|
||
|
|
|
||
|
|
else:
|
||
|
|
# Otherwise, expect only a single WRTE
|
||
|
|
cmd, data = conn.ReadUntil(b'WRTE')
|
||
|
|
|
||
|
|
# WRTE cmd from device will follow with stdout data
|
||
|
|
stdout_stream.write(data)
|
||
|
|
|
||
|
|
else:
|
||
|
|
|
||
|
|
# No cmd provided means we should just expect a single line from the terminal. Use this sparingly
|
||
|
|
cmd, data = conn.ReadUntil(b'WRTE')
|
||
|
|
if cmd == b'WRTE':
|
||
|
|
# WRTE cmd from device will follow with stdout data
|
||
|
|
stdout_stream.write(data)
|
||
|
|
else:
|
||
|
|
print("Unhandled cmd: {}".format(cmd))
|
||
|
|
|
||
|
|
cleaned_stdout_stream = BytesIO()
|
||
|
|
if clean_stdout:
|
||
|
|
stdout_bytes = stdout_stream.getvalue()
|
||
|
|
|
||
|
|
bsruns = {} # Backspace runs tracking
|
||
|
|
next_start_pos = 0
|
||
|
|
last_run_pos, last_run_len = find_backspace_runs(stdout_bytes, next_start_pos)
|
||
|
|
|
||
|
|
if last_run_pos != -1 and last_run_len != 0:
|
||
|
|
bsruns.update({last_run_pos: last_run_len})
|
||
|
|
cleaned_stdout_stream.write(stdout_bytes[next_start_pos:(last_run_pos - last_run_len)])
|
||
|
|
next_start_pos += last_run_pos + last_run_len
|
||
|
|
|
||
|
|
while last_run_pos != -1:
|
||
|
|
last_run_pos, last_run_len = find_backspace_runs(stdout_bytes[next_start_pos:], next_start_pos)
|
||
|
|
|
||
|
|
if last_run_pos != -1:
|
||
|
|
bsruns.update({last_run_pos: last_run_len})
|
||
|
|
cleaned_stdout_stream.write(stdout_bytes[next_start_pos:(last_run_pos - last_run_len)])
|
||
|
|
next_start_pos += last_run_pos + last_run_len
|
||
|
|
|
||
|
|
cleaned_stdout_stream.write(stdout_bytes[next_start_pos:])
|
||
|
|
|
||
|
|
else:
|
||
|
|
cleaned_stdout_stream.write(stdout_stream.getvalue())
|
||
|
|
|
||
|
|
stdout = cleaned_stdout_stream.getvalue()
|
||
|
|
|
||
|
|
# Strip original cmd that will come back in stdout
|
||
|
|
if original_cmd and strip_cmd:
|
||
|
|
findstr = original_cmd.encode('utf-8') + b'\r\r\n'
|
||
|
|
pos = stdout.find(findstr)
|
||
|
|
while pos >= 0:
|
||
|
|
stdout = stdout.replace(findstr, b'')
|
||
|
|
pos = stdout.find(findstr)
|
||
|
|
|
||
|
|
if b'\r\r\n' in stdout:
|
||
|
|
stdout = stdout.split(b'\r\r\n')[1]
|
||
|
|
|
||
|
|
# Strip delim if requested
|
||
|
|
# TODO: Handling stripping partial delims here - not a deal breaker the way we're handling it now
|
||
|
|
if delim and strip_delim:
|
||
|
|
stdout = stdout.replace(delim, b'')
|
||
|
|
|
||
|
|
stdout = stdout.rstrip()
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
print("InteractiveShell exception (most likely timeout): {}".format(e))
|
||
|
|
|
||
|
|
return stdout
|