Module pycti.arbinspoofer.arbin_spoofer
Expand source code
import socket
import threading
import struct
import copy
from pycti.messages import Msg, MessageABC
class ChannelData:
__chan_readings_list = []
__chan_readings_lock = threading.Lock()
def __init__(self, num_channels):
"""
Container class that will hold all of the specific channel data for ArbinSpoofer.
Parameters
----------
num_channels : int
Number of channels in our hypothetical Arbin cycler.
"""
self.num_channels = num_channels
# Create channel_readings for all of the channels.
for i in range(0, self.num_channels):
channel_readings = {}
for key, item in Msg.ChannelInfo.Server.msg_specific_template.items():
channel_readings[key] = copy.deepcopy(item['value'])
channel_readings['channel'] = i
with self.__chan_readings_lock:
self.__chan_readings_list.append(
copy.deepcopy(channel_readings))
def fetch_channel_readings(self, channel) -> dict:
"""
Returns the status message for a specified channel.
Parameters
----------
channel : int
The channel to return the status for.
Returns
-------
status : dict
The status message for the requested channel. Empty if channel larger than number of channel
"""
if channel > self.num_channels:
return {}
else:
with self.__chan_readings_lock:
return copy.deepcopy(self.__chan_readings_list[channel])
def update_channel_readings(self, channel, updated_readings):
"""
Updates the stored channel readings for the specified channel.
Parameters
----------
channel : int
The channel to update the readings for.
updated_status : dict
Complete or partial dictionary of status values to update.
Returns
-------
success : bool
Returns True if all values in the updated_status were used to update the channel_status_array.
"""
if channel > self.num_channels:
return False
else:
for key in updated_readings.keys():
if key not in Msg.ChannelInfo.Server.msg_specific_template.keys():
return False
with self.__chan_readings_lock:
for key in updated_readings.keys():
self.__chan_readings_list[channel][key] = updated_readings[key]
class SocketWorker:
"""
Generic worker class that will respond to client socket requests.
Default setup as an echo server. Child classes should overwrite the
the `_process_client_msg()` method with their own responses.
"""
__receive_msg_timeout_s = 0.5
__msg_buffer_size_bytes = 2**12
__stop_lock = threading.Lock()
__stop = False
def __init__(self, s: socket.socket, channel_data: ChannelData):
"""
Creates the thread to service client requests.
Parameters
----------
s : socket.socket
Socket connection to client.
"""
self.__channel_data = channel_data
self.stop = False
self.__client_thread = threading.Thread(
target=self.__service_loop,
args=(s,),
daemon=True
)
self.__client_thread.start()
def __service_loop(self, s: socket.socket):
"""
Forever loop to service client requests. Wait to receive a message. If no messages is
received before the timeout then check to see if stop command has been issued. Loop is
also broken if client breaks connection by sending b''.
Parameters
----------
s : socket.socket
Socket connection to client.
"""
s.settimeout(self.__receive_msg_timeout_s)
rx_msg_length_format = MessageABC.base_template['msg_length']['format']
rx_msg_length_start_byte = MessageABC.base_template['msg_length']['start_byte']
rx_msg_length_end_byte = MessageABC.base_template['msg_length']['start_byte'] + struct.calcsize(
rx_msg_length_format)
while True:
try:
rx_msg = s.recv(self.__msg_buffer_size_bytes)
if not rx_msg:
break
# Keep reading message in pieces until rx_msg is as long as expected_rx_msg_len
expected_rx_msg_len = struct.unpack(
rx_msg_length_format, rx_msg[rx_msg_length_start_byte:rx_msg_length_end_byte])[0]
while len(rx_msg) < expected_rx_msg_len:
rx_msg += s.recv(self.__msg_buffer_size_bytes)
tx_msg = self.__process_client_msg(rx_msg)
s.sendall(tx_msg)
except socket.timeout:
with self.__stop_lock:
if self.__stop:
break
s.close()
def __process_client_msg(self, rx_msg):
"""
Takes the incoming client message and generates a response.
Parameters
----------
rx_msg : bytearray
The client message received.
Returns
-------
tx_msg : PyBytesObject
The client response.
"""
# Determine command code to sort message
cmd_code_format = MessageABC.base_template['command_code']['format']
cmd_code_start_byte = MessageABC.base_template['command_code']['start_byte']
cmd_code_end_byte = MessageABC.base_template['command_code']['start_byte'] + + struct.calcsize(
cmd_code_format)
cmd_code = struct.unpack(
cmd_code_format, rx_msg[cmd_code_start_byte:cmd_code_end_byte])[0]
if cmd_code == Msg.Login.Client.command_code:
rx_msg_dict = Msg.Login.Client.unpack(rx_msg)
tx_msg = Msg.Login.Server.pack({'num_channels':self.__channel_data.num_channels})
elif cmd_code == Msg.ChannelInfo.Client.command_code:
rx_msg_dict = Msg.ChannelInfo.Client.unpack(rx_msg)
channel_values = self.__channel_data.fetch_channel_readings(
rx_msg_dict['channel'])
tx_msg = Msg.ChannelInfo.Server.pack(channel_values)
elif cmd_code == Msg.AssignSchedule.Client.command_code:
rx_msg_dict = Msg.AssignSchedule.Client.unpack(rx_msg)
tx_msg = Msg.AssignSchedule.Server.pack(
{'channel': rx_msg_dict['channel']})
elif cmd_code == Msg.StartSchedule.Client.command_code:
rx_msg_dict = Msg.StartSchedule.Client.unpack(rx_msg)
tx_msg = Msg.StartSchedule.Server.pack(
{'channel': rx_msg_dict['channel']})
elif cmd_code == Msg.StopSchedule.Client.command_code:
rx_msg_dict = Msg.StopSchedule.Client.unpack(rx_msg)
tx_msg = Msg.StopSchedule.Server.pack(
{'channel': rx_msg_dict['channel']})
elif cmd_code == Msg.SetMetaVariable.Client.command_code:
rx_msg_dict = Msg.SetMetaVariable.Client.unpack(rx_msg)
tx_msg = Msg.SetMetaVariable.Server.pack(
{'channel': rx_msg_dict['channel']})
else:
tx_msg = bytearray([])
return tx_msg
def is_alive(self):
"""
Method to call to see if the client service thread is still running.
Returns
-------
running : bool
True of False based on whether or not the client thread is running.
"""
return self.__client_thread.is_alive()
def kill_worker(self):
"""
Method to stop client service loop.
"""
if self.__client_thread.is_alive():
with self.__stop_lock:
self.__stop = True
self.__client_thread.join()
class ArbinSpoofer:
__client_connect_timeout_s = 0.5
__stop_servers_lock = threading.Lock()
__stop_servers = False
def __init__(self, config: dict):
"""
Class to mimic behavior of Arbin cycler MITSPro control server. The class is currently dumb
and just sends back basic channel messages without any notion of channel status.
It could be expanded in future.
Parameters
----------
config : dict
A configuration dictionary containing the following fields:
`ip`: The server IP address to host from. Most often will be 'localhost' for testing.
`port`: The port to use for the server.
`num_channels`: The number of channel our fictitious cycler has.
"""
self.__channel_data = ChannelData(config['num_channels'])
self.__server_thread = threading.Thread(
target=self.__server_loop,
args=(config, SocketWorker,),
daemon=True
)
def start(self):
"""
Starts the server loops.
"""
self.__server_thread.start()
def update_channel_status(self, channel, updated_readings):
"""
Updates the stored channel status for the specified channel.
Parameters
----------
channel : int
The channel to update the status for.
updated_readings : dict
Complete or partial dictionary of status values to update.
Returns
-------
success : bool
Returns True if all values in the updated_status were used to update the channel_status_array.
"""
return self.__channel_data.update_channel_readings(channel, updated_readings)
def __server_loop(self, sock_config: dict, Worker: SocketWorker):
"""
Creates a server and forever loop to service client socket requests.
Parameters
----------
sock_config : dict
A configuration for the socket containing the IP and port number.
Worker : SocketWorker
A reference to the worker class that will service individual client connections.
"""
# List that will hold all the workers to service client connections.
client_workers = []
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind((sock_config["ip"], sock_config["port"]))
sock.settimeout(self.__client_connect_timeout_s)
sock.listen()
while True:
try:
client_connection = sock.accept()[0]
client_workers.append(
Worker(client_connection, self.__channel_data))
except socket.timeout:
with self.__stop_servers_lock:
# If stop command is issued then kill all workers.
if self.__stop_servers:
for worker in client_workers:
if worker.is_alive():
worker.kill_worker()
break
# Remove any workers that made have died from disconnecting clients.
client_workers[:] = [
worker for worker in client_workers if worker.is_alive()
]
def stop(self):
"""
Stop the server loops.
"""
with self.__stop_servers_lock:
self.__stop_servers = True
self.__server_thread.join()
def __del__(self):
self.stop()
Classes
class ArbinSpoofer (config: dict)
-
Class to mimic behavior of Arbin cycler MITSPro control server. The class is currently dumb and just sends back basic channel messages without any notion of channel status. It could be expanded in future.
Parameters
config
:dict
-
A configuration dictionary containing the following fields:
ip
: The server IP address to host from. Most often will be 'localhost' for testing.port
: The port to use for the server.num_channels
: The number of channel our fictitious cycler has.
Expand source code
class ArbinSpoofer: __client_connect_timeout_s = 0.5 __stop_servers_lock = threading.Lock() __stop_servers = False def __init__(self, config: dict): """ Class to mimic behavior of Arbin cycler MITSPro control server. The class is currently dumb and just sends back basic channel messages without any notion of channel status. It could be expanded in future. Parameters ---------- config : dict A configuration dictionary containing the following fields: `ip`: The server IP address to host from. Most often will be 'localhost' for testing. `port`: The port to use for the server. `num_channels`: The number of channel our fictitious cycler has. """ self.__channel_data = ChannelData(config['num_channels']) self.__server_thread = threading.Thread( target=self.__server_loop, args=(config, SocketWorker,), daemon=True ) def start(self): """ Starts the server loops. """ self.__server_thread.start() def update_channel_status(self, channel, updated_readings): """ Updates the stored channel status for the specified channel. Parameters ---------- channel : int The channel to update the status for. updated_readings : dict Complete or partial dictionary of status values to update. Returns ------- success : bool Returns True if all values in the updated_status were used to update the channel_status_array. """ return self.__channel_data.update_channel_readings(channel, updated_readings) def __server_loop(self, sock_config: dict, Worker: SocketWorker): """ Creates a server and forever loop to service client socket requests. Parameters ---------- sock_config : dict A configuration for the socket containing the IP and port number. Worker : SocketWorker A reference to the worker class that will service individual client connections. """ # List that will hold all the workers to service client connections. client_workers = [] sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) sock.bind((sock_config["ip"], sock_config["port"])) sock.settimeout(self.__client_connect_timeout_s) sock.listen() while True: try: client_connection = sock.accept()[0] client_workers.append( Worker(client_connection, self.__channel_data)) except socket.timeout: with self.__stop_servers_lock: # If stop command is issued then kill all workers. if self.__stop_servers: for worker in client_workers: if worker.is_alive(): worker.kill_worker() break # Remove any workers that made have died from disconnecting clients. client_workers[:] = [ worker for worker in client_workers if worker.is_alive() ] def stop(self): """ Stop the server loops. """ with self.__stop_servers_lock: self.__stop_servers = True self.__server_thread.join() def __del__(self): self.stop()
Methods
def start(self)
-
Starts the server loops.
Expand source code
def start(self): """ Starts the server loops. """ self.__server_thread.start()
def stop(self)
-
Stop the server loops.
Expand source code
def stop(self): """ Stop the server loops. """ with self.__stop_servers_lock: self.__stop_servers = True self.__server_thread.join()
def update_channel_status(self, channel, updated_readings)
-
Updates the stored channel status for the specified channel.
Parameters
channel
:int
- The channel to update the status for.
updated_readings
:dict
- Complete or partial dictionary of status values to update.
Returns
success
:bool
- Returns True if all values in the updated_status were used to update the channel_status_array.
Expand source code
def update_channel_status(self, channel, updated_readings): """ Updates the stored channel status for the specified channel. Parameters ---------- channel : int The channel to update the status for. updated_readings : dict Complete or partial dictionary of status values to update. Returns ------- success : bool Returns True if all values in the updated_status were used to update the channel_status_array. """ return self.__channel_data.update_channel_readings(channel, updated_readings)
class ChannelData (num_channels)
-
Container class that will hold all of the specific channel data for ArbinSpoofer.
Parameters
num_channels : int Number of channels in our hypothetical Arbin cycler.
Expand source code
class ChannelData: __chan_readings_list = [] __chan_readings_lock = threading.Lock() def __init__(self, num_channels): """ Container class that will hold all of the specific channel data for ArbinSpoofer. Parameters ---------- num_channels : int Number of channels in our hypothetical Arbin cycler. """ self.num_channels = num_channels # Create channel_readings for all of the channels. for i in range(0, self.num_channels): channel_readings = {} for key, item in Msg.ChannelInfo.Server.msg_specific_template.items(): channel_readings[key] = copy.deepcopy(item['value']) channel_readings['channel'] = i with self.__chan_readings_lock: self.__chan_readings_list.append( copy.deepcopy(channel_readings)) def fetch_channel_readings(self, channel) -> dict: """ Returns the status message for a specified channel. Parameters ---------- channel : int The channel to return the status for. Returns ------- status : dict The status message for the requested channel. Empty if channel larger than number of channel """ if channel > self.num_channels: return {} else: with self.__chan_readings_lock: return copy.deepcopy(self.__chan_readings_list[channel]) def update_channel_readings(self, channel, updated_readings): """ Updates the stored channel readings for the specified channel. Parameters ---------- channel : int The channel to update the readings for. updated_status : dict Complete or partial dictionary of status values to update. Returns ------- success : bool Returns True if all values in the updated_status were used to update the channel_status_array. """ if channel > self.num_channels: return False else: for key in updated_readings.keys(): if key not in Msg.ChannelInfo.Server.msg_specific_template.keys(): return False with self.__chan_readings_lock: for key in updated_readings.keys(): self.__chan_readings_list[channel][key] = updated_readings[key]
Methods
def fetch_channel_readings(self, channel) ‑> dict
-
Returns the status message for a specified channel.
Parameters
channel
:int
- The channel to return the status for.
Returns
status
:dict
- The status message for the requested channel. Empty if channel larger than number of channel
Expand source code
def fetch_channel_readings(self, channel) -> dict: """ Returns the status message for a specified channel. Parameters ---------- channel : int The channel to return the status for. Returns ------- status : dict The status message for the requested channel. Empty if channel larger than number of channel """ if channel > self.num_channels: return {} else: with self.__chan_readings_lock: return copy.deepcopy(self.__chan_readings_list[channel])
def update_channel_readings(self, channel, updated_readings)
-
Updates the stored channel readings for the specified channel.
Parameters
channel
:int
- The channel to update the readings for.
updated_status
:dict
- Complete or partial dictionary of status values to update.
Returns
success
:bool
- Returns True if all values in the updated_status were used to update the channel_status_array.
Expand source code
def update_channel_readings(self, channel, updated_readings): """ Updates the stored channel readings for the specified channel. Parameters ---------- channel : int The channel to update the readings for. updated_status : dict Complete or partial dictionary of status values to update. Returns ------- success : bool Returns True if all values in the updated_status were used to update the channel_status_array. """ if channel > self.num_channels: return False else: for key in updated_readings.keys(): if key not in Msg.ChannelInfo.Server.msg_specific_template.keys(): return False with self.__chan_readings_lock: for key in updated_readings.keys(): self.__chan_readings_list[channel][key] = updated_readings[key]
class SocketWorker (s: socket.socket, channel_data: ChannelData)
-
Generic worker class that will respond to client socket requests. Default setup as an echo server. Child classes should overwrite the the
_process_client_msg()
method with their own responses.Creates the thread to service client requests.
Parameters
s
:socket.socket
- Socket connection to client.
Expand source code
class SocketWorker: """ Generic worker class that will respond to client socket requests. Default setup as an echo server. Child classes should overwrite the the `_process_client_msg()` method with their own responses. """ __receive_msg_timeout_s = 0.5 __msg_buffer_size_bytes = 2**12 __stop_lock = threading.Lock() __stop = False def __init__(self, s: socket.socket, channel_data: ChannelData): """ Creates the thread to service client requests. Parameters ---------- s : socket.socket Socket connection to client. """ self.__channel_data = channel_data self.stop = False self.__client_thread = threading.Thread( target=self.__service_loop, args=(s,), daemon=True ) self.__client_thread.start() def __service_loop(self, s: socket.socket): """ Forever loop to service client requests. Wait to receive a message. If no messages is received before the timeout then check to see if stop command has been issued. Loop is also broken if client breaks connection by sending b''. Parameters ---------- s : socket.socket Socket connection to client. """ s.settimeout(self.__receive_msg_timeout_s) rx_msg_length_format = MessageABC.base_template['msg_length']['format'] rx_msg_length_start_byte = MessageABC.base_template['msg_length']['start_byte'] rx_msg_length_end_byte = MessageABC.base_template['msg_length']['start_byte'] + struct.calcsize( rx_msg_length_format) while True: try: rx_msg = s.recv(self.__msg_buffer_size_bytes) if not rx_msg: break # Keep reading message in pieces until rx_msg is as long as expected_rx_msg_len expected_rx_msg_len = struct.unpack( rx_msg_length_format, rx_msg[rx_msg_length_start_byte:rx_msg_length_end_byte])[0] while len(rx_msg) < expected_rx_msg_len: rx_msg += s.recv(self.__msg_buffer_size_bytes) tx_msg = self.__process_client_msg(rx_msg) s.sendall(tx_msg) except socket.timeout: with self.__stop_lock: if self.__stop: break s.close() def __process_client_msg(self, rx_msg): """ Takes the incoming client message and generates a response. Parameters ---------- rx_msg : bytearray The client message received. Returns ------- tx_msg : PyBytesObject The client response. """ # Determine command code to sort message cmd_code_format = MessageABC.base_template['command_code']['format'] cmd_code_start_byte = MessageABC.base_template['command_code']['start_byte'] cmd_code_end_byte = MessageABC.base_template['command_code']['start_byte'] + + struct.calcsize( cmd_code_format) cmd_code = struct.unpack( cmd_code_format, rx_msg[cmd_code_start_byte:cmd_code_end_byte])[0] if cmd_code == Msg.Login.Client.command_code: rx_msg_dict = Msg.Login.Client.unpack(rx_msg) tx_msg = Msg.Login.Server.pack({'num_channels':self.__channel_data.num_channels}) elif cmd_code == Msg.ChannelInfo.Client.command_code: rx_msg_dict = Msg.ChannelInfo.Client.unpack(rx_msg) channel_values = self.__channel_data.fetch_channel_readings( rx_msg_dict['channel']) tx_msg = Msg.ChannelInfo.Server.pack(channel_values) elif cmd_code == Msg.AssignSchedule.Client.command_code: rx_msg_dict = Msg.AssignSchedule.Client.unpack(rx_msg) tx_msg = Msg.AssignSchedule.Server.pack( {'channel': rx_msg_dict['channel']}) elif cmd_code == Msg.StartSchedule.Client.command_code: rx_msg_dict = Msg.StartSchedule.Client.unpack(rx_msg) tx_msg = Msg.StartSchedule.Server.pack( {'channel': rx_msg_dict['channel']}) elif cmd_code == Msg.StopSchedule.Client.command_code: rx_msg_dict = Msg.StopSchedule.Client.unpack(rx_msg) tx_msg = Msg.StopSchedule.Server.pack( {'channel': rx_msg_dict['channel']}) elif cmd_code == Msg.SetMetaVariable.Client.command_code: rx_msg_dict = Msg.SetMetaVariable.Client.unpack(rx_msg) tx_msg = Msg.SetMetaVariable.Server.pack( {'channel': rx_msg_dict['channel']}) else: tx_msg = bytearray([]) return tx_msg def is_alive(self): """ Method to call to see if the client service thread is still running. Returns ------- running : bool True of False based on whether or not the client thread is running. """ return self.__client_thread.is_alive() def kill_worker(self): """ Method to stop client service loop. """ if self.__client_thread.is_alive(): with self.__stop_lock: self.__stop = True self.__client_thread.join()
Methods
def is_alive(self)
-
Method to call to see if the client service thread is still running.
Returns
running
:bool
- True of False based on whether or not the client thread is running.
Expand source code
def is_alive(self): """ Method to call to see if the client service thread is still running. Returns ------- running : bool True of False based on whether or not the client thread is running. """ return self.__client_thread.is_alive()
def kill_worker(self)
-
Method to stop client service loop.
Expand source code
def kill_worker(self): """ Method to stop client service loop. """ if self.__client_thread.is_alive(): with self.__stop_lock: self.__stop = True self.__client_thread.join()