class Client():
def __init__(self):
self.current_protocol = "SEP" # changed to FMP later
self.expected_incoming_msg_type = "RES_INIT"
self.ver = 1
self.NET_PATH = "./network"
self.FOLDER_PATH = "./client"
if (self.NET_PATH[-1] != '/') and (self.NET_PATH[-1] != '\\'):
self.NET_PATH += '/'
if (self.FOLDER_PATH[-1] != '/') and (self.FOLDER_PATH[-1] != '\\'):
self.FOLDER_PATH += '/'
self.DST_ADDR = None
self.OWN_ADDR = input('Type your own address: ') # Ex. A, B, or C
self.netif = network_interface(self.NET_PATH, self.OWN_ADDR)
self.protocols = Protocols()
self.session_key = None
self.client_prikey = None
self.client_pubkey = None
self.server_pubkey = None
self.timeout = 30
def clean_keys(self):
self.client_prikey = None
self.client_pubkey = None
self.server_pubkey = None
def load_client_prikey(self):
try:
key = open(
self.FOLDER_PATH + "private_keys/" + self.OWN_ADDR + ".pem",
"rb").read()
self.client_prikey = RSA.import_key(key)
except:
print("Error: could not load client private key")
def load_client_pubkey(self):
try:
key = open(
self.FOLDER_PATH + "public_keys/" + self.OWN_ADDR + ".pem",
"rb").read()
self.client_pubkey = RSA.import_key(key)
except:
print("Error: could not load client public key")
def load_server_pubkey(self):
try:
key = open(
self.FOLDER_PATH + "public_keys/" + self.DST_ADDR + ".pem",
"rb").read()
self.server_pubkey = RSA.import_key(key)
except:
print("Error: could not load server public key")
def fix_length(self, input_byte, desired_length):
if (len(input_byte) > desired_length):
print("Error fixing the length of byte")
return None
res = b""
for i in range(desired_length - len(input_byte)):
res += b"\x00"
res += input_byte
return res
def handle_options(self):
try:
opts, args = getopt.getopt(sys.argv[1:],
shortopts='hp:a',
longopts=['help', 'path=', 'addr='])
except getopt.GetoptError:
print('wrong options')
sys.exit(1)
def generate_nonce(self):
return secrets.token_bytes(32)
def SEP_ENC(self, ptext, key):
cipher_rsa = PKCS1_OAEP.new(key)
ctext = cipher_rsa.encrypt(ptext)
return ctext
def SEP_DEC(self, ctext, key):
cipher_rsa = PKCS1_OAEP.new(key)
ptext = cipher_rsa.decrypt(ctext)
return ptext
def SEP_SIGN(self, message, key):
h = SHA256.new(message)
signature = pkcs1_15.new(key).sign(h)
return signature
def SEP_GEN(self, ver, type_sep, type_msg, sender_id, recipient_id,
encrypted_content, sig_key):
msg = b""
msg += ver
msg += type_sep
msg += type_msg
msg += len(encrypted_content).to_bytes(length=5,
byteorder='big') #msg_len
msg += sender_id
msg += recipient_id
msg += encrypted_content
sig = self.SEP_SIGN(msg, sig_key)
#print("Generated sig: ")
#print(sig)
msg += sig
return msg
def SEP_REQ_INIT(self):
# generate SEP REQ_INIT
#header
ver = self.ver.to_bytes(length=1, byteorder='big')
type_sep = self.protocols.protocol2num("SEP").to_bytes(length=1,
byteorder='big')
type_msg = self.protocols.sep_type2num("REQ_INIT").to_bytes(
length=1, byteorder='big')
sender_id = self.fix_length(self.OWN_ADDR.encode('utf-8'), 4)
recipient_id = self.fix_length(self.DST_ADDR.encode('utf-8'), 4)
#keys to be used
#server_public_key = RSA.import_key(open("public_keys/"+self.DST_ADDR+".pem").read())
#client_private_key = RSA.import_key(open("private_keys/"+self.OWN_ADDR+".pem").read())
#content
nonce = self.generate_nonce()
#pubkey_a = self.client_pubkey
pubkey_a = self.client_pubkey.export_key()
#with open("public_keys/"+self.OWN_ADDR+".pem", "rb") as f: pubkey_a = f.read()
ctext = b""
ctext += len(pubkey_a).to_bytes(length=4, byteorder='big')
ctext += pubkey_a
nonce_encrypted = self.SEP_ENC(nonce, self.server_pubkey)
ctext += nonce_encrypted
#send msg
msg = self.SEP_GEN(ver, type_sep, type_msg, sender_id, recipient_id,
ctext, self.client_prikey)
self.netif.send_msg(self.DST_ADDR, msg)
#print("Generated nonce: ")
#print(nonce)
#print("encrypted nonce: ")
#print(nonce_encrypted)
return nonce
def SEP_REQ_LOGIN(self):
password = getpass('Type your login password: '******'big')
type_sep = self.protocols.protocol2num("SEP").to_bytes(length=1,
byteorder='big')
type_msg = self.protocols.sep_type2num("REQ_LOGIN").to_bytes(
length=1, byteorder='big')
sender_id = self.fix_length(self.OWN_ADDR.encode('utf-8'), 4)
recipient_id = self.fix_length(self.DST_ADDR.encode('utf-8'), 4)
#keys to be used
#server_public_key = RSA.import_key(open("public_keys/"+self.DST_ADDR+".pem").read())
#client_private_key = RSA.import_key(open("private_keys/"+self.OWN_ADDR+".pem").read())
#content
ctext = self.SEP_ENC(password.encode('utf-8'), self.server_pubkey)
#send msg
msg = self.SEP_GEN(ver, type_sep, type_msg, sender_id, recipient_id,
ctext, self.client_prikey)
self.netif.send_msg(self.DST_ADDR, msg)
return True
def read_SEP(self, msg):
type_msg = msg[2]
if type_msg == self.protocols.sep_type2num(
self.expected_incoming_msg_type):
print("(OK): message type matches")
else:
print("Error: unexpected message type")
return None
msg_len = int.from_bytes(msg[3:8], 'big') #int
if len(msg) == 16 + msg_len + 256:
print("(OK): message length matches")
else:
print("Error: message length does not match")
return None
sender_id = msg[8:12].decode('utf-8').strip(chr(0)) #string
recipient_id = msg[12:16].decode('utf-8').strip(chr(0)) #string
if recipient_id == self.OWN_ADDR:
print("(OK): correct recipient address")
else:
print("Error: wrong recipient address")
return None
# verify signature
sig = msg[-256:]
h = SHA256.new(msg[:-256])
try:
pkcs1_15.new(self.server_pubkey).verify(h, sig)
print("(OK): signature valid")
except:
print("Error: signature invalid")
return None
if type_msg == self.protocols.sep_type2num('RES_INIT'):
#client_private_key = RSA.import_key(open("private_keys/"+self.OWN_ADDR+".pem").read())
#server_public_key = RSA.import_key(open("public_keys/"+self.DST_ADDR+".pem").read())
encrypted_nonce = msg[16:-256]
# decrypt nonce
cipher_rsa = PKCS1_OAEP.new(self.client_prikey)
decrypted_nonce = cipher_rsa.decrypt(encrypted_nonce)
return decrypted_nonce
elif type_msg == self.protocols.sep_type2num('RES_LOGIN'):
encrypted_msg_content = msg[16:-256]
# decrypt
cipher_rsa = PKCS1_OAEP.new(self.client_prikey)
decrypted_msg_content = cipher_rsa.decrypt(encrypted_msg_content)
login_status = decrypted_msg_content[0]
if login_status == 0:
print("Error: login failed due to incorrect password")
return None
elif login_status == 1: # login sucess
session_key = decrypted_msg_content[1:]
print("(OK): password correct")
return session_key
else:
print("Error: uknown response type")
else:
print("Error: unknown message type")
return None
def connect(self):
self.DST_ADDR = input('Type a server address: ')
self.load_client_prikey()
self.load_client_pubkey()
self.load_server_pubkey()
# send SEP REQ_INIT
nonce = None
while True:
nonce = self.SEP_REQ_INIT()
if nonce is not None:
print("SEP REQ INIT sent")
break
# wait for SEP RES_INIT
self.expected_incoming_msg_type = "RES_INIT"
stime = time.time()
while True:
status, msg = self.netif.receive_msg(blocking=False)
if status:
print("SEP RES_INIT received")
nonce_returned = self.read_SEP(msg)
if nonce == nonce_returned:
print("(OK): server verified")
print("SEP RES_INIT accpeted")
break
else:
print("Error: server verification failed")
return None
#check request timeout
etime = time.time()
if (etime - stime > self.timeout): # 20 sec timeout
print("Connection request timed out:", self.timeout, "sec")
return None
time.sleep(0.5)
# send SEP REQ_LOGIN
while True:
res = self.SEP_REQ_LOGIN()
if res:
break
# wait for SEP RES_LOGIN
self.expected_incoming_msg_type = "RES_LOGIN"
stime = time.time()
while True:
status, msg = self.netif.receive_msg(blocking=False)
if status:
print("SEP RES_LOGIN received")
session_key = self.read_SEP(msg)
if session_key is not None:
return session_key
else:
continue
etime = time.time()
if (etime - stime > self.timeout):
print("Connection request timed out:", self.timeout, "sec")
return None
time.sleep(0.5)
def start(self):
session_key = None
while True:
session_key = self.connect()
if session_key is not None:
break
print("Connection established with", self.DST_ADDR)
print(session_key)
class ScoreCardView:
def __init__(self):
# Initialize Data Objects
self.dvh = None
self.dvh_counts = []
self.dvh_counts_for_plot = None
self.bin_count = 0
self.bin_counts = None
self.roi_keys = None
self.roi_names = None
self.roi_key_map = None
self.plans = None
self.structures = None
self.protocol_data = None
self.roi_override = {}
self.aliases = StructureAliases()
self.protocols = Protocols()
self.source_data = ColumnDataSource(data=dict(roi_name=[], roi_template=[], roi_key=[], volume=[], min_dose=[],
mean_dose=[], max_dose=[], constraint=[], constraint_calc=[],
pass_fail=[], calc_type=[]))
self.source_plot = ColumnDataSource(data=dict(x=[], y=[], color=[], roi=[], roi_key=[]))
self.colors = itertools.cycle(palette)
self.__define_layout_objects()
self.__do_bind()
self.__do_layout()
self.update_protocol_data()
self.initialize_source_data()
def __define_layout_objects(self):
# Report heading data
self.select_plan = Select(title='Plan:', width=400)
self.button_refresh_plans = Button(label='Scan DICOM Inbox', button_type='primary')
self.select_protocol = Select(title='Protocol:', options=self.protocols.protocol_names, value='TG101', width=150)
self.select_fx = Select(title='Fractions:', value='3', options=self.fractionation_options, width=60)
self.button_calculate = Button(label='Calculate Scorecard', button_type='primary')
self.button_delete_roi = Button(label='Delete Constraint', button_type='warning')
self.button_calculate_dvhs = Button(label='Calculate DVHs', button_type='primary', width=200)
self.select_roi_template = Select(title='Template ROI:')
self.select_roi = Select(title='Plan ROI:')
self.max_dose_volume = Div(text="<b>Point defined as %scc" % MAX_DOSE_VOLUME)
self.columns = [TableColumn(field="roi_template", title="Template ROI"),
TableColumn(field="roi_name", title="ROI"),
TableColumn(field='volume', title='Volume (cc)', formatter=NumberFormatter(format="0.00")),
TableColumn(field='min_dose', title='Min Dose (Gy)', formatter=NumberFormatter(format="0.00")),
TableColumn(field='mean_dose', title='Mean Dose (Gy)', formatter=NumberFormatter(format="0.00")),
TableColumn(field='max_dose', title='Max Dose (Gy)', formatter=NumberFormatter(format="0.00")),
TableColumn(field='constraint', title='Constraint'),
TableColumn(field='constraint_calc', title='Value', formatter=NumberFormatter(format="0.00")),
TableColumn(field='pass_fail', title='Pass/Fail', formatter=self.__pass_fail_formatter)]
self.data_table = DataTable(source=self.source_data, columns=self.columns, index_position=None,
width=1000, height=300)
tools = "pan,wheel_zoom,box_zoom,reset,crosshair,save"
self.plot = figure(plot_width=800, plot_height=475, tools=tools, active_drag="box_zoom")
# Set x and y axis labels
self.plot.xaxis.axis_label = "Dose (Gy)"
self.plot.yaxis.axis_label = "Normalized Volume"
self.plot.min_border_left = 60
self.plot.min_border_bottom = 60
self.plot.add_tools(HoverTool(show_arrow=False, line_policy='next',
tooltips=[('Plan ROI', '@roi'),
('Dose', '$x'),
('Volume', '$y')]))
self.plot.xaxis.axis_label_text_font_size = "12pt"
self.plot.yaxis.axis_label_text_font_size = "12pt"
self.plot.xaxis.major_label_text_font_size = "10pt"
self.plot.yaxis.major_label_text_font_size = "10pt"
self.plot.yaxis.axis_label_text_baseline = "bottom"
self.plot.lod_factor = 100 # level of detail during interactive plot events
self.plot.multi_line('x', 'y', source=self.source_plot,
selection_color='color', line_width=3, alpha=0,
line_dash='solid', nonselection_alpha=0, selection_alpha=1)
columns = [TableColumn(field="roi", title="Select Structures to Plot")]
self.plot_rois = DataTable(source=self.source_plot, columns=columns, index_position=None,
width=200, height=350)
def __do_bind(self):
self.button_calculate.on_click(self.initialize_source_data)
self.button_delete_roi.on_click(self.delete_selected_rows)
self.select_protocol.on_change('value', self.protocol_listener)
self.select_fx.on_change('value', self.fx_listener)
self.button_refresh_plans.on_click(self.update_plan_options)
self.select_plan.on_change('value', self.plan_listener)
self.select_roi_template.on_change('value', self.template_roi_listener)
self.select_roi.on_change('value', self.roi_listener)
self.source_data.selected.on_change('indices', self.source_select)
self.button_calculate_dvhs.on_click(self.update_dvh_plot)
def __do_layout(self):
self.layout = column(self.button_refresh_plans,
row(self.select_plan, self.select_protocol, self.select_fx),
row(self.button_calculate, self.button_delete_roi),
row(self.select_roi_template, self.select_roi),
self.max_dose_volume,
self.data_table,
row(self.plot, Spacer(width=10), column(self.button_calculate_dvhs,
self.plot_rois)))
@property
def __pass_fail_formatter(self):
# Data tables
# custom js to highlight mismatches in red with white text
template = """
<div style="background:<%=
(function colorfrommismatch(){
if(pass_fail != "" ){
if(pass_fail == "Fail"){
return('HIGHLIGHT_COLOR_FAIL')
}
if(pass_fail == "Pass"){
return('HIGHLIGHT_COLOR_PASS')
}
}
}()) %>;
color: <%=
(function colorfrommismatch(){
if(pass_fail == "Fail"){return('TEXT_COLOR_FAIL')}
}()) %>;">
<%= value %>
</div>
"""
template = template.replace('HIGHLIGHT_COLOR_FAIL', 'red')
template = template.replace('TEXT_COLOR_FAIL', 'white')
template = template.replace('HIGHLIGHT_COLOR_PASS', 'lightgreen')
template = template.replace('TEXT_COLOR_PASS', 'black')
return HTMLTemplateFormatter(template=template)
# Properties -------------------------------------------------------------------
@property
def fractionation(self):
return "%sFx" % self.select_fx.value
@property
def fractionation_options(self):
return self.protocols.get_fractionations(self.protocol)
@property
def protocol(self):
return self.select_protocol.value
@property
def current_struct_file(self):
return self.plans[self.select_plan.value]['rtstruct']
# Listeners -------------------------------------------------------------------
def protocol_listener(self, attr, old, new):
self.select_fx.options = self.fractionation_options
if self.select_fx.value not in self.select_fx.options:
self.select_fx.value = self.select_fx.options[0]
else: # Changing select_fx.value will prompt the following two lines
self.update_protocol_data()
self.initialize_source_data()
def fx_listener(self, attr, old, new):
self.update_protocol_data()
self.initialize_source_data()
def plan_listener(self, attr, old, new):
self.roi_override = {}
if new in list(self.plans):
self.initialize_source_data()
def template_roi_listener(self, attr, old, new):
self.update_roi_select()
def roi_listener(self, attr, old, new):
template_rois = self.source_data.data['roi_template']
indices = [i for i, roi in enumerate(template_rois) if roi == self.select_roi_template.value]
if indices:
patches = {'roi_name': [(i, new) for i in indices]}
self.source_data.patch(patches)
if new:
self.roi_override[self.select_roi_template.value] = new
self.button_calculate.button_type = 'success'
total = len(indices)
for i in indices:
self.button_calculate.label = 'Calculating ScoreCard %s of %s' % (i+1, total)
key = self.roi_key_map[new]
self.calculate_dvh(key)
self.update_table_row(i, key)
self.source_data.patch({'roi_key': [(i, key)]})
self.update_constraint(i)
self.button_calculate.label = 'Calculate Scorecard'
self.button_calculate.button_type = 'primary'
else:
if self.button_calculate.button_type == 'primary' and self.select_roi_template.value in self.roi_override:
self.roi_override.pop(self.select_roi_template.value)
patches = {'volume': [(i, 0.) for i in indices],
'min_dose': [(i, 0.) for i in indices],
'mean_dose': [(i, 0.) for i in indices],
'max_dose': [(i, 0.) for i in indices],
'constraint_calc': [(i, 0.) for i in indices],
'pass_fail': [(i, '') for i in indices],
'roi_key': [(i, '') for i in indices]}
self.source_data.patch(patches)
def source_select(self, attr, old, new):
if new:
self.select_roi_template.value = self.source_data.data['roi_template'][new[0]]
# Methods -------------------------------------------------------------------
def update_protocol_data(self):
self.protocol_data = self.protocols.get_column_data(self.protocol, self.fractionation)
def initialize_source_data(self):
data = self.protocol_data
self.bin_counts = None
self.bin_count = None
self.dvh = {}
self.dvh_counts_for_plot = []
self.dvh_counts = []
row_count = len(data['roi_template'])
new_data = {'roi_template': data['roi_template'],
'roi_key': [''] * row_count,
'roi_name': [''] * row_count,
'volume': [''] * row_count,
'min_dose': [''] * row_count,
'mean_dose': [''] * row_count,
'max_dose': [''] * row_count,
'constraint': data['string_rep'],
'constraint_calc': [''] * row_count,
'pass_fail': [''] * row_count,
'calc_type': data['calc_type']}
self.source_plot.data = {'x': [], 'y': [], 'color': [], 'roi': [], 'roi_key': []}
self.source_data.data = new_data
self.update_roi_template_select()
if self.select_plan.value:
self.button_calculate.label = 'Calculating Scorecard...'
self.button_calculate.button_type = 'success'
self.update_plan_structures()
def delete_selected_rows(self):
selected_indices = self.source_data.selected.indices
selected_indices.sort(reverse=True)
if not selected_indices:
selected_indices = [0]
data = self.source_data.data
for index in selected_indices:
for key in list(data):
data[key].pop(index)
self.source_data.data = data
self.source_data.selected.indices = []
def update_roi_template_select(self):
options = list(set(self.source_data.data['roi_template']))
options.sort()
self.select_roi_template.options = options
if self.select_roi_template.value not in options:
self.select_roi_template.value = options[0]
def update_plan_options(self):
self.button_refresh_plans.button_type = 'success'
self.button_refresh_plans.label = 'Updating...'
self.plans = get_plans(INBOX_DIR)
self.select_plan.options = list(self.plans)
self.button_refresh_plans.button_type = 'primary'
self.button_refresh_plans.label = 'Scan DICOM Inbox'
if self.select_plan.value not in list(self.plans):
self.select_plan.value = list(self.plans)[0]
def update_plan_structures(self):
self.structures = dicomparser.DicomParser(self.current_struct_file).GetStructures()
self.roi_keys = [key for key in self.structures if self.structures[key]['type'].upper() != 'MARKER']
self.roi_names = [str(self.structures[key]['name']) for key in self.roi_keys]
self.roi_key_map = {name: self.roi_keys[i] for i, name in enumerate(self.roi_names)}
self.select_roi.options = [''] + self.roi_names
self.update_roi_select()
self.match_rois()
def update_roi_select(self):
index = self.source_data.data['roi_template'].index(self.select_roi_template.value)
self.select_roi.value = self.source_data.data['roi_name'][index]
def match_rois(self):
matches = self.aliases.match_protocol_rois(self.source_data.data['roi_template'], self.roi_names)
patches = {'roi_name': [], 'roi_key': []}
for i, protocol_roi in enumerate(self.source_data.data['roi_template']):
if protocol_roi in list(matches):
match = matches[protocol_roi]
else:
match = ''
if protocol_roi in list(self.roi_override):
match = self.roi_override[protocol_roi]
if match:
key = self.roi_key_map[match]
else:
key = ''
patches['roi_name'].append((i, match))
patches['roi_key'].append((i, key))
self.source_data.patch(patches)
self.update_roi_select()
self.calculate_protocol_dvhs()
def update_table_row(self, index, key):
patch = {'volume': [(index, self.dvh[key].volume)],
'min_dose': [(index, self.dvh[key].min)],
'mean_dose': [(index, self.dvh[key].mean)],
'max_dose': [(index, self.dvh[key].max)]}
self.source_data.patch(patch)
def calculate_dvh(self, key):
if key not in list(self.dvh):
files = self.plans[self.select_plan.value]
self.dvh[key] = dvhcalc.get_dvh(files['rtstruct'], files['rtdose'], key)
def calculate_dvhs(self):
current_button_type = self.button_calculate.button_type
current_button_label = self.button_calculate.label
self.button_calculate_dvhs.button_type = 'success'
files = self.plans[self.select_plan.value]
total = len(self.roi_keys)
for i, key in enumerate(list(self.roi_keys)):
self.button_calculate_dvhs.label = 'Calculating DVH: %s of %s' % (i+1, total)
if key not in list(self.dvh):
self.dvh[key] = dvhcalc.get_dvh(files['rtstruct'], files['rtdose'], key)
self.dvh_counts = [self.dvh[key].counts for key in self.roi_keys]
self.button_calculate.button_type = current_button_type
self.button_calculate.label = current_button_label
def calculate_protocol_dvhs(self):
self.dvh = {}
total = len([key for key in self.source_data.data['roi_key'] if key != ''])
counter = 1
for i, key in enumerate(self.source_data.data['roi_key']):
if key:
self.button_calculate.label = 'Calculating Scorecard %s of %s' % (counter, total)
self.calculate_dvh(key)
self.update_table_row(i, key)
self.update_constraint(i)
counter += 1
self.button_calculate.label = 'Calculate Scorecard'
self.button_calculate.button_type = 'primary'
def update_constraint(self, index):
if self.source_data.data['roi_name'][index]:
constraint = self.calculate_constraint(index)
operator = self.protocol_data['operator'][index]
threshold = self.protocol_data['threshold_value'][index]
if constraint is not None:
if operator == '<':
status = constraint < threshold
else:
status = constraint > threshold
status = ['Fail', 'Pass'][status]
else:
status = ''
self.source_data.patch({'constraint_calc': [(index, constraint)],
'pass_fail': [(index, status)]})
def calculate_constraint(self, index):
dvh = self.dvh[self.source_data.data['roi_key'][index]]
calc_type = self.source_data.data['calc_type'][index]
input_value = self.protocol_data['input_value'][index]
if calc_type == 'Volume':
ans = dvh.dose_constraint(input_value, volume_units='cm3')
return float(str(ans).split(' ')[0])
if calc_type == 'Dose':
ans = dvh.volume_constraint(input_value, dose_units='Gy')
return float(str(ans).split(' ')[0])
if calc_type == 'Mean':
return self.source_data.data['mean_dose'][index]
if calc_type == 'MVS':
ans = dvh.volume_constraint(input_value, dose_units='Gy')
ans = float(str(ans).split(' ')[0])
return self.source_data.data['volume'][index] - ans
return None
@property
def volumes(self):
return np.array([self.dvh[key].volume for key in self.roi_keys])
def pad_dvh_counts(self):
self.bin_count = max([len(dvh) for dvh in self.dvh_counts])
self.dvh_counts_for_plot = []
self.bin_counts = []
for dvh in self.dvh_counts:
if dvh[0]:
self.dvh_counts_for_plot.append(np.divide(dvh, dvh[0]))
else:
self.dvh_counts_for_plot.append(np.array([0] * self.bin_count))
np.append(self.dvh_counts_for_plot[-1], [0]) # ensure last value is 0
x_axis = np.divide(np.array(list(range(len(self.dvh_counts_for_plot[-1])))), 100.)
self.bin_counts.append(x_axis)
def update_dvh_plot(self):
self.calculate_dvhs()
self.pad_dvh_counts()
colors = [color for j, color in zip(range(len(self.dvh_counts_for_plot)), self.colors)]
self.source_plot.data = {'x': self.bin_counts,
'y': self.dvh_counts_for_plot,
'color': colors,
'roi': self.roi_names,
'roi_key': self.roi_keys}
self.button_calculate_dvhs.label = 'Calculate DVHs'
self.button_calculate_dvhs.button_type = 'primary'
def __init__(self):
logging.debug(
f"{self.__class__.__name__ } - Initialising master server.")
self.transport = None
self.storage = Storage()
self.protocols = Protocols()
class MasterServer:
def __init__(self):
logging.debug(
f"{self.__class__.__name__ } - Initialising master server.")
self.transport = None
self.storage = Storage()
self.protocols = Protocols()
def connection_made(self, transport) -> NoReturn:
self.transport = transport
def datagram_received(self, data: bytes, address: Tuple[str,
int]) -> NoReturn:
response = None
logging.debug(
f"{self.__class__.__name__ } - Recieved {data} from {address}")
result: Dict = self.protocols.parse_data(data)
if result.get("class", None) == "B2M":
response = self.handle_client(result)
elif result.get("class", None) == "S2M":
response = self.handle_server(result, address)
else:
pass
if response:
self.send_response(response, address)
def send_response(self, response: bytes, address: Tuple[str,
int]) -> NoReturn:
logging.debug(
f"{self.__class__.__name__ } - Sending {response} to {address}")
self.transport.sendto(response, address)
def handle_client(self, result: Dict) -> bytes:
logging.debug(f"{self.__class__.__name__ } - Header belongs to client")
response_header = result.get("resp", None)
server_list = self.storage.list_server_addresses(result.get("game"))
processed_server_list = [self.pack_address(_) for _ in server_list]
return self.create_response(response_header, processed_server_list)
def handle_server(self, result, address: Tuple[str, int]) -> bytes:
logging.debug(f"{self.__class__.__name__ } - Header belongs to server")
server = GameServer(address, result)
if self.storage.get_server(server):
self.storage.update_server(server)
else:
self.storage.create_server(server)
if not server.active:
self.storage.server_shutdown(server)
return result.get("resp", None)
@staticmethod
def create_response(header: bytes, response: List[bytes]) -> bytes:
seperator = b""
if header:
response.insert(0, header)
return seperator.join(response)
@staticmethod
def pack_address(address: str) -> bytes:
"""
Takes string formatted address;
eg, '192.168.0.1:27910'
Converts to 6 byte binary string.
H = unsigned short
"""
port_format = ">H"
ip, port = address.split(":")
ip = ip_address(ip).packed
port = struct.pack(port_format, int(port))
return ip + port
class Server():
def __init__(self):
self.ver = 1
self.current_protocol = 'SEP' # changed to FMP later
self.expected_incoming_msg_type = "REQ_INIT"
self.OWN_ADDR = "Z" #input("Type your own address: ")
print("Starting server as Z...")
self.DST_ADDR = None
self.NET_PATH = "./network"
self.FOLDER_PATH = "./server"
if (self.NET_PATH[-1] != '/') and (self.NET_PATH[-1] != '\\'):
self.NET_PATH += '/'
if (self.FOLDER_PATH[-1] != '/') and (self.FOLDER_PATH[-1] != '\\'):
self.FOLDER_PATH += '/'
self.netif = network_interface(self.NET_PATH, self.OWN_ADDR)
self.protocols = Protocols()
self.session_key = None
self.client_pubkey = None
self.client_pubkey_len = None
self.server_prikey = None
#to decrypted saved password and prikey
self.passphrase = "sepfmp"
self.iv = b'\x8c\xbcW\xcf\x0b\xa6\x00\xec\xa7\x94\xd2\x9a\x01Z\xd7\xfc'
self.timeout = 30
def load_server_prikey(self):
try:
key = open(
self.FOLDER_PATH + "private_keys/" + self.OWN_ADDR + ".pem",
"rb").read()
self.server_prikey = RSA.import_key(key,
passphrase=self.passphrase)
except:
print("Error: could not load server private key")
def fix_length(self, input_byte, desired_length):
if (len(input_byte) > desired_length):
print("Error fixing the length of byte")
return
res = b""
for i in range(desired_length - len(input_byte)):
res += b"\x00"
res += input_byte
return res
def generate_session_key(self):
return secrets.token_bytes(32)
def mkpad(self, s, size):
s = s.encode("utf-8") # UTF-8文字列をバイト列に変換する
pad = b' ' * (size - len(s) % size) # 特定の長さの倍数にするための空白を生成
return s + pad
def verify_password(self, password):
key = self.mkpad(self.passphrase, 16)
cipher = AES.new(key, AES.MODE_CBC, self.iv)
ctext = open(self.FOLDER_PATH + "password/" + self.DST_ADDR,
"rb").read()
correct_pwd = cipher.decrypt(ctext)
correct_pwd = correct_pwd.strip(b" ")
if correct_pwd == password:
return True
return False
def SEP_ENC(self, ptext, key):
cipher_rsa = PKCS1_OAEP.new(key)
ctext = cipher_rsa.encrypt(ptext)
return ctext
def SEP_DEC(self, ctext, key):
cipher_rsa = PKCS1_OAEP.new(key)
ptext = cipher_rsa.decrypt(ctext)
return ptext
def SEP_SIGN(self, message, key):
h = SHA256.new(message)
signature = pkcs1_15.new(key).sign(h)
return signature
def SEP_GEN(self, ver, type_sep, type_msg, sender_id, recipient_id,
encrypted_content, sig_key):
msg = b""
msg += ver
msg += type_sep
msg += type_msg
msg += len(encrypted_content).to_bytes(length=5,
byteorder='big') #msg_len
msg += sender_id
msg += recipient_id
msg += encrypted_content
sig = self.SEP_SIGN(msg, sig_key)
#print("Generated sig: ")
#print(sig)
msg += sig
return msg
def SEP_RES_INIT(self, nonce):
#header
ver = self.ver.to_bytes(length=1, byteorder='big')
type_sep = self.protocols.protocol2num("SEP").to_bytes(length=1,
byteorder='big')
type_msg = self.protocols.sep_type2num("RES_INIT").to_bytes(
length=1, byteorder='big')
sender_id = self.fix_length(self.OWN_ADDR.encode('utf-8'), 4)
recipient_id = self.fix_length(self.DST_ADDR.encode('utf-8'), 4)
#content - encrypted nonce
ctext = self.SEP_ENC(nonce, self.client_pubkey)
#send msg
msg = self.SEP_GEN(ver, type_sep, type_msg, sender_id, recipient_id,
ctext, self.server_prikey)
self.netif.send_msg(self.DST_ADDR, msg)
def SEP_RES_LOGIN(self, status, session_key=None):
#header
ver = self.ver.to_bytes(length=1, byteorder='big')
type_sep = self.protocols.protocol2num("SEP").to_bytes(length=1,
byteorder='big')
type_msg = self.protocols.sep_type2num("RES_LOGIN").to_bytes(
length=1, byteorder='big')
sender_id = self.fix_length(self.OWN_ADDR.encode('utf-8'), 4)
recipient_id = self.fix_length(self.DST_ADDR.encode('utf-8'), 4)
login_status = None
ctext = None
if status:
s = 1
login_status = s.to_bytes(length=1, byteorder='big')
ptext = login_status + session_key
ctext = self.SEP_ENC(ptext, self.client_pubkey)
else:
s = 0
login_status = s.to_bytes(length=1, byteorder='big')
ctext = self.SEP_ENC(login_status, self.client_pubkey)
msg = self.SEP_GEN(ver, type_sep, type_msg, sender_id, recipient_id,
ctext, self.server_prikey)
self.netif.send_msg(self.DST_ADDR, msg)
def read_SEP(self, msg):
type_msg = msg[2]
if type_msg == self.protocols.sep_type2num(
self.expected_incoming_msg_type):
print("(OK): protocol type matches")
else:
print("Error: unexpected message type")
return None
msg_len = int.from_bytes(msg[3:8], 'big') #int
if len(msg) == 16 + msg_len + 256:
print("(OK): message length matches")
else:
print("Error: message length does not match")
return None
sender_id = msg[8:12].decode('utf-8').strip(chr(0)) #string
recipient_id = msg[12:16].decode('utf-8').strip(chr(0)) #string
if recipient_id == self.OWN_ADDR:
print("(OK): correct recipient address")
else:
print("Error: wrong recipient address")
return None
if self.client_pubkey is None and type_msg == self.protocols.sep_type2num(
'REQ_INIT'):
self.client_pubkey_len = int.from_bytes(msg[16:20], 'big')
self.client_pubkey = RSA.import_key(msg[20:20 +
self.client_pubkey_len])
#verify signature
sig = msg[-256:]
h = SHA256.new(msg[:-256])
try:
pkcs1_15.new(self.client_pubkey).verify(h, sig)
print("(OK): sginature valid")
except:
print("Error: signature invalid")
return
if type_msg == self.protocols.sep_type2num('REQ_INIT'):
encrypted_nonce = msg[20 + self.client_pubkey_len:-256]
#decrypt nonce
#self.server_prikey = RSA.import_key(open("private_keys/"+self.OWN_ADDR+".pem").read())
cipher_rsa = PKCS1_OAEP.new(self.server_prikey)
decrypted_nonce = cipher_rsa.decrypt(encrypted_nonce)
#print("type_msg", type_msg)
#print("msg_len",msg_len)
#print("sender",sender_id)
#print("recipient", recipient_id)
#print("pubkey_a")
#print(pubkey_a)
#print("encrypted nonce")
#print(encrypted_nonce)
#print("sig")
#print(sig)
#update internal parameters
self.DST_ADDR = sender_id
return decrypted_nonce
elif type_msg == self.protocols.sep_type2num('REQ_LOGIN'):
encrypted_password = msg[16:-256]
#decrypt password
#self.server_prikey = RSA.import_key(open("private_keys/"+self.OWN_ADDR+".pem").read())
cipher_rsa = PKCS1_OAEP.new(self.server_prikey)
decrypted_password = cipher_rsa.decrypt(encrypted_password)
return decrypted_password
else:
print("Error: unknown message type")
return None
def read_msg(self, msg):
ver = msg[0]
type_protocol = msg[1]
print(ver)
print(type_protocol)
if type_protocol != self.protocols.protocol2num(self.current_protocol):
print(
"Error: unexpected protocol type - received message protocol="
+ type_protocol)
return None, None, None
if type_protocol == self.protocols.protocol2num("SEP"):
return ver, type_protocol, self.read_SEP(msg)
elif type_protocol == self.protocols.protocol2num("FMP"):
pass
#Todo: create self.read_FMP
def connect(self):
self.current_protocol = "SEP"
self.expected_incoming_msg_type = "REQ_INIT"
status = None
msg = None
# wait for SEP REQ_INIT
self.expected_incoming_msg_type = "REQ_INIT"
while True:
status, msg = self.netif.receive_msg(blocking=True)
print("SEP REQ_INIT received")
ver, type_protocol, decrypted_nonce = self.read_msg(msg)
if decrypted_nonce is not None:
print("SEP REQ_INIT accepted")
self.SEP_RES_INIT(decrypted_nonce)
print("SEP RES_INIT sent")
break
else:
print("Error: SEP REQ_INIT rejected")
return None
# wait for SEP REQ_LOGIN
self.expected_incoming_msg_type = "REQ_LOGIN"
stime = time.time()
while True:
status, msg = self.netif.receive_msg(blocking=False)
if status:
print("SEP REQ_LOGIN received")
ver, type_protocol, password = self.read_msg(msg)
if password is not None:
print("SEP REQ_LOGIN accepted")
# TODO SEP RES LOGIN
if self.verify_password(password):
print("(OK): password ok")
session_key = self.generate_session_key()
self.SEP_RES_LOGIN(True, session_key=session_key)
print("SEP RES_LOGIN sent")
return session_key
else:
print("(!!): password incorrect")
self.SEP_RES_LOGIN(False)
print("SEP RES_LOGIN sent")
continue # keep waiting for correct password
#cehck request timeout
etime = time.time()
if (etime - stime > self.timeout):
print("Connection request timed out:", self.timeout, "sec")
return None
time.sleep(0.5)
def start(self):
self.load_server_prikey()
session_key = None
while True:
session_key = self.connect()
if session_key is not None:
break
print("Connection established with", self.DST_ADDR)
class Client():
def __init__(self):
self.current_protocol = "SEP" # changed to FMP later
self.expected_incoming_msg_type = "RES_INIT"
self.ver = 1
self.NET_PATH = "./network"
self.FOLDER_PATH = "./client"
if (self.NET_PATH[-1] != '/') and (self.NET_PATH[-1] != '\\'):
self.NET_PATH += '/'
if (self.FOLDER_PATH[-1] != '/') and (self.FOLDER_PATH[-1] != '\\'):
self.FOLDER_PATH += '/'
self.DST_ADDR = None
self.OWN_ADDR = input('Type your own address: ') # Ex. A, B, or C
self.netif = network_interface(self.NET_PATH, self.OWN_ADDR)
self.protocols = Protocols()
self.timeout = 30
self.sleep_time = 0.1
self.session_key = None
self.client_prikey = None
self.client_pubkey = None
self.server_pubkey = None
#fmp seq num
self.seq_last_received = 0
self.seq_last_sent = 0
self.DATA_LOC = "my_data"
def clean_keys(self):
self.client_prikey = None
self.client_pubkey = None
self.server_pubkey = None
def load_client_prikey(self):
try:
key = open(
self.FOLDER_PATH + "private_keys/" + self.OWN_ADDR + ".pem",
"rb").read()
self.client_prikey = RSA.import_key(key)
except:
print("Error: could not load client private key")
def load_client_pubkey(self):
try:
key = open(
self.FOLDER_PATH + "public_keys/" + self.OWN_ADDR + ".pem",
"rb").read()
self.client_pubkey = RSA.import_key(key)
except:
print("Error: could not load client public key")
def load_server_pubkey(self):
try:
key = open(
self.FOLDER_PATH + "public_keys/" + self.DST_ADDR + ".pem",
"rb").read()
self.server_pubkey = RSA.import_key(key)
except:
print("Error: could not load server public key")
def fix_length(self, input_byte, desired_length):
if (len(input_byte) > desired_length):
print("Error fixing the length of byte")
return None
res = b""
for i in range(desired_length - len(input_byte)):
res += b"\x00"
res += input_byte
return res
def handle_options(self):
try:
opts, args = getopt.getopt(sys.argv[1:],
shortopts='hp:a',
longopts=['help', 'path=', 'addr='])
except getopt.GetoptError:
print('wrong options')
sys.exit(1)
def generate_nonce(self):
return secrets.token_bytes(32)
def SEP_ENC(self, ptext, key):
cipher_rsa = PKCS1_OAEP.new(key)
ctext = cipher_rsa.encrypt(ptext)
return ctext
def SEP_DEC(self, ctext, key):
cipher_rsa = PKCS1_OAEP.new(key)
ptext = cipher_rsa.decrypt(ctext)
return ptext
def SEP_SIGN(self, message, key):
h = SHA256.new(message)
signature = pkcs1_15.new(key).sign(h)
return signature
def SEP_GEN(self, ver, type_sep, type_msg, sender_id, recipient_id,
encrypted_content, sig_key):
msg = b""
msg += ver
msg += type_sep
msg += type_msg
msg += len(encrypted_content).to_bytes(length=5,
byteorder='big') #msg_len
msg += sender_id
msg += recipient_id
msg += encrypted_content
sig = self.SEP_SIGN(msg, sig_key)
#print("Generated sig: ")
#print(sig)
msg += sig
return msg
def SEP_REQ_INIT(self):
# generate SEP REQ_INIT
#header
ver = self.ver.to_bytes(length=1, byteorder='big')
type_sep = self.protocols.protocol2num("SEP").to_bytes(length=1,
byteorder='big')
type_msg = self.protocols.sep_type2num("REQ_INIT").to_bytes(
length=1, byteorder='big')
sender_id = self.fix_length(self.OWN_ADDR.encode('utf-8'), 4)
recipient_id = self.fix_length(self.DST_ADDR.encode('utf-8'), 4)
#content
nonce = self.generate_nonce()
#pubkey_a = self.client_pubkey
pubkey_a = self.client_pubkey.export_key()
#with open("public_keys/"+self.OWN_ADDR+".pem", "rb") as f: pubkey_a = f.read()
ctext = b""
ctext += len(pubkey_a).to_bytes(length=4, byteorder='big')
ctext += pubkey_a
nonce_encrypted = self.SEP_ENC(nonce, self.server_pubkey)
ctext += nonce_encrypted
#send msg
msg = self.SEP_GEN(ver, type_sep, type_msg, sender_id, recipient_id,
ctext, self.client_prikey)
self.netif.send_msg(self.DST_ADDR, msg)
#print("Generated nonce: ")
#print(nonce)
#print("encrypted nonce: ")
#print(nonce_encrypted)
return nonce
def SEP_REQ_LOGIN(self):
password = getpass('Type your login password: '******'big')
type_sep = self.protocols.protocol2num("SEP").to_bytes(length=1,
byteorder='big')
type_msg = self.protocols.sep_type2num("REQ_LOGIN").to_bytes(
length=1, byteorder='big')
sender_id = self.fix_length(self.OWN_ADDR.encode('utf-8'), 4)
recipient_id = self.fix_length(self.DST_ADDR.encode('utf-8'), 4)
#keys to be used
#server_public_key = RSA.import_key(open("public_keys/"+self.DST_ADDR+".pem").read())
#client_private_key = RSA.import_key(open("private_keys/"+self.OWN_ADDR+".pem").read())
#content
ctext = self.SEP_ENC(password.encode('utf-8'), self.server_pubkey)
#send msg
msg = self.SEP_GEN(ver, type_sep, type_msg, sender_id, recipient_id,
ctext, self.client_prikey)
self.netif.send_msg(self.DST_ADDR, msg)
return True
def read_SEP(self, msg):
type_msg = msg[2]
if type_msg == self.protocols.sep_type2num(
self.expected_incoming_msg_type):
print("(OK): message type matches")
else:
print("Error: unexpected message type")
return None
msg_len = int.from_bytes(msg[3:8], 'big') #int
if len(msg) == 16 + msg_len + 256:
print("(OK): message length matches")
else:
print("Error: message length does not match")
return None
sender_id = msg[8:12].decode('utf-8').strip(chr(0)) #string
recipient_id = msg[12:16].decode('utf-8').strip(chr(0)) #string
if recipient_id == self.OWN_ADDR:
print("(OK): correct recipient address")
else:
print("Error: wrong recipient address")
return None
# verify signature
sig = msg[-256:]
h = SHA256.new(msg[:-256])
try:
pkcs1_15.new(self.server_pubkey).verify(h, sig)
print("(OK): signature valid")
except:
print("Error: signature invalid")
return None
if type_msg == self.protocols.sep_type2num('RES_INIT'):
#client_private_key = RSA.import_key(open("private_keys/"+self.OWN_ADDR+".pem").read())
#server_public_key = RSA.import_key(open("public_keys/"+self.DST_ADDR+".pem").read())
encrypted_nonce = msg[16:-256]
# decrypt nonce
cipher_rsa = PKCS1_OAEP.new(self.client_prikey)
decrypted_nonce = cipher_rsa.decrypt(encrypted_nonce)
return decrypted_nonce
elif type_msg == self.protocols.sep_type2num('RES_LOGIN'):
encrypted_msg_content = msg[16:-256]
# decrypt
cipher_rsa = PKCS1_OAEP.new(self.client_prikey)
decrypted_msg_content = cipher_rsa.decrypt(encrypted_msg_content)
login_status = decrypted_msg_content[0]
if login_status == 0:
print("Error: login failed due to incorrect password")
return None
elif login_status == 1: # login sucess
session_key = decrypted_msg_content[1:]
print("(OK): password correct")
return session_key
else:
print("Error: uknown response type")
else:
print("Error: unknown message type")
return None
def read_msg(self, msg):
ver = msg[0]
type_protocol = msg[1]
if type_protocol != self.protocols.protocol2num(self.current_protocol):
print(
"Error: unexpected protocol type - received message protocol="
+ type_protocol)
return None, None, None
if type_protocol == self.protocols.protocol2num("SEP"):
return ver, type_protocol, self.read_SEP(msg)
elif type_protocol == self.protocols.protocol2num("FMP"):
return ver, type_protocol, self.read_FMP(msg)
def connect(self):
self.DST_ADDR = input('Type a server address: ')
self.load_client_prikey()
self.load_client_pubkey()
self.load_server_pubkey()
# send SEP REQ_INIT
nonce = None
while True:
nonce = self.SEP_REQ_INIT()
if nonce is not None:
print("SEP REQ INIT sent")
break
# wait for SEP RES_INIT
self.expected_incoming_msg_type = "RES_INIT"
stime = time.time()
while True:
status, msg = self.netif.receive_msg(blocking=False)
if status:
print("SEP RES_INIT received")
ver, type_protocol, nonce_returned = self.read_msg(msg)
if nonce == nonce_returned:
print("(OK): server verified")
print("SEP RES_INIT accpeted")
break
else:
print("Error: server verification failed")
return None
#check request timeout
etime = time.time()
if (etime - stime > self.timeout): # 20 sec timeout
print("Connection request timed out:", self.timeout, "sec")
return None
time.sleep(self.sleep_time)
while True:
# send SEP REQ_LOGIN
while True:
res = self.SEP_REQ_LOGIN()
if res:
break
# wait for SEP RES_LOGIN
self.expected_incoming_msg_type = "RES_LOGIN"
stime = time.time()
while True:
status, msg = self.netif.receive_msg(blocking=False)
if status:
print("SEP RES_LOGIN received")
ver, type_protocol, session_key = self.read_msg(msg)
if session_key is not None:
return session_key
else:
break
etime = time.time()
if (etime - stime > self.timeout):
print("Connection request timed out:", self.timeout, "sec")
return None
time.sleep(self.sleep_time)
def gen_iv(self):
return secrets.token_bytes(16)
def FMP_ENC(self, ptext):
iv = self.gen_iv()
len(self.session_key)
cipher = AES.new(self.session_key, AES.MODE_CBC, iv=iv)
ctext = cipher.encrypt(pad(ptext, AES.block_size))
return iv, ctext
def FMP_MAC(self, msg):
h = HMAC.new(self.session_key, digestmod=SHA256)
h.update(msg)
return h.digest()
def FMP_DEC(self, iv, ctext):
cipher = AES.new(self.session_key, AES.MODE_CBC, iv=iv)
ptext = unpad(cipher.decrypt(ctext), AES.block_size)
return ptext
def FMP_GEN(self, type_msg, content):
ver = self.ver.to_bytes(length=1, byteorder='big')
type_fmp = self.protocols.protocol2num("FMP").to_bytes(length=1,
byteorder='big')
type_msg = self.protocols.fmp_type2num(type_msg).to_bytes(
length=1, byteorder='big')
sender_id = self.fix_length(self.OWN_ADDR.encode('utf-8'), 4)
recipient_id = self.fix_length(self.DST_ADDR.encode('utf-8'), 4)
iv, encrypted_content = self.FMP_ENC(content)
msg = b""
msg += ver
msg += type_fmp
msg += type_msg
msg += (len(iv) + len(encrypted_content)).to_bytes(length=5,
byteorder=('big'))
msg += sender_id
msg += recipient_id
msg += iv
msg += encrypted_content
mac = self.FMP_MAC(msg)
msg += mac
return msg
def FMP_LISTEN(self, command):
# wait for FMP RES
self.expected_incoming_msg_type = command
stime = time.time()
while True:
status, msg = self.netif.receive_msg(blocking=False)
if status:
print("FMP_" + command, "response received")
ver, type_protocol, content = self.read_msg(msg)
if content is not None:
return content
else:
print("Error: response message invalid")
continue
etime = time.time()
if (etime - stime > self.timeout):
print("Error: no response received")
return None
time.sleep(self.sleep_time)
return None
def FMP_SEND(self, msg):
self.netif.send_msg(self.DST_ADDR, msg)
self.seq_last_sent += 1
def FMP_MKD(self):
dir_name = input("Type directory name: ")
# content
# 2 bytes for seq num
content = (self.seq_last_sent + 1).to_bytes(2, 'big')
content += dir_name.encode('utf-8')
msg = self.FMP_GEN("MKD", content)
self.FMP_SEND(msg)
print("FMP_MKD request sent")
res = self.FMP_LISTEN("MKD")
res = int.from_bytes(res, 'big')
if res == 0:
print("Error: MKD failed")
elif res == 1:
print("MKD: directory created")
def FMP_RMD(self):
dir_name = input("Type directory name: ")
confirm = input("All sub contents will be deleted: (y/n) ")
if confirm != "y":
print("RMD cancelled")
return
content = (self.seq_last_sent + 1).to_bytes(2, 'big')
content += dir_name.encode('utf-8')
msg = self.FMP_GEN("RMD", content)
self.FMP_SEND(msg)
print("FMP_RMD request sent")
res = self.FMP_LISTEN("RMD")
res = int.from_bytes(res, 'big')
if res == 0:
print("Error: RMD failed")
elif res == 1:
print("RMD: directory deleted")
def FMP_GWD(self):
content = (self.seq_last_sent + 1).to_bytes(2, 'big')
msg = self.FMP_GEN("GWD", content)
self.FMP_SEND(msg)
print("FMP_GWD request sent")
res = self.FMP_LISTEN("GWD")
if res is not None:
res = res.decode('utf-8')
print(res)
def FMP_CWD(self):
dir_name = input("Type destination: ")
# content
# 2 bytes for seq num
content = (self.seq_last_sent + 1).to_bytes(2, 'big')
content += dir_name.encode('utf-8')
msg = self.FMP_GEN("CWD", content)
self.FMP_SEND(msg)
print("FMP_CWD request sent")
res = self.FMP_LISTEN("CWD")
res = int.from_bytes(res, 'big')
if res == 0:
print("Error: CWD failed")
elif res == 1:
print("CWD: current directory changed")
def FMP_LST(self):
content = (self.seq_last_sent + 1).to_bytes(2, 'big')
msg = self.FMP_GEN("LST", content)
self.FMP_SEND(msg)
print("FMP_LST request sent")
res = self.FMP_LISTEN("LST")
if res is not None:
res = res.decode('utf-8')
res = list(res.split(","))
for i in res:
print(i)
def FMP_UPL(self):
file_name = input("Type file name:")
try:
f_cont = open(
os.path.join(self.FOLDER_PATH, self.DATA_LOC, file_name),
"rb").read()
except:
print("Error: failed to read the file")
return
# content
# 2 bytes for seq num
content = (self.seq_last_sent + 1).to_bytes(2, 'big')
content += len(file_name).to_bytes(2, ('big'))
content += file_name.encode('utf-8')
content += f_cont
msg = self.FMP_GEN("UPL", content)
self.FMP_SEND(msg)
print("FMP_UPL request sent")
res = self.FMP_LISTEN("UPL")
res = int.from_bytes(res, 'big')
if res == 0:
print("Error: UPL failed")
elif res == 1:
print("UPL: file uploaded")
def FMP_DNL(self):
file_name = input("Type file name: ")
content = (self.seq_last_sent + 1).to_bytes(2, 'big')
content += file_name.encode('utf-8')
msg = self.FMP_GEN("DNL", content)
self.FMP_SEND(msg)
print("FMP_DNL request sent")
res = self.FMP_LISTEN("DNL")
fcont = res[1:].decode('utf-8')
res = res[0]
if res == 0:
print("Error: DNL failed")
elif res == 1:
open(os.path.join(self.FOLDER_PATH, self.DATA_LOC, file_name),
"w").write(fcont)
print("DNL: file downloaded")
def FMP_RMF(self):
file_name = input("Type file name: ")
content = (self.seq_last_sent + 1).to_bytes(2, 'big')
content += file_name.encode('utf-8')
msg = self.FMP_GEN("RMF", content)
self.FMP_SEND(msg)
print("FMP_RMF request sent")
res = self.FMP_LISTEN("RMF")
res = int.from_bytes(res, 'big')
if res == 0:
print("Error: RMF failed")
elif res == 1:
print("RMF: file removed")
def FMP_END(self):
confirm = input("Disconnect and terminate session: (y/n)")
if confirm != 'y':
print("END cancelled")
return
content = (self.seq_last_sent + 1).to_bytes(2, 'big')
msg = self.FMP_GEN("END", content)
self.FMP_SEND(msg)
print("FMP_END request sent")
res = self.FMP_LISTEN("END")
res = int.from_bytes(res, 'big')
if res == 0:
print("Error: END failed")
elif res == 1:
os.sys.exit(0)
print("END: session terminated")
def read_FMP(self, msg):
type_command = msg[2]
msg_len = int.from_bytes(msg[3:8], 'big') #int
if len(msg) == 16 + msg_len + 32:
print("(OK): message length matches")
else:
print("Error: message length does not match")
return None
sender_id = msg[8:12].decode('utf-8').strip(chr(0)) #string
if sender_id == self.DST_ADDR:
print("(OK): correct sender address")
else:
print("Error: wrong sender address")
recipient_id = msg[12:16].decode('utf-8').strip(chr(0)) #string
if recipient_id == self.OWN_ADDR:
print("(OK): correct recipient address")
else:
print("Error: wrong recipient address")
return None
iv = msg[16:32]
ctext = msg[32:-32]
mac = msg[-32:]
# verify mac
h = HMAC.new(self.session_key, digestmod=SHA256)
h.update(msg[:-32])
try:
h.verify(mac)
print("(OK): MAC ok")
except ValueError:
print("Error: MAC verification failed")
return None
# decrypt
ptext = self.FMP_DEC(iv, ctext)
seq_num = int.from_bytes(ptext[0:2], byteorder='big')
content = ptext[2:]
# check seq_num
if seq_num <= self.seq_last_received:
print("Error: invalid sequence number")
return None
else:
self.seq_last_received = seq_num
print("(OK): sequence number ok")
return content
def fmp_process(self):
command = input("Type command: ")
if command == "MKD":
self.FMP_MKD()
elif command == "RMD":
self.FMP_RMD()
elif command == "GWD":
self.FMP_GWD()
elif command == "CWD":
self.FMP_CWD()
elif command == "LST":
self.FMP_LST()
elif command == "UPL":
self.FMP_UPL()
elif command == "DNL":
self.FMP_DNL()
elif command == "RMF":
self.FMP_RMF()
elif command == "END":
self.FMP_END()
else:
print("Error: unknown command")
def start(self):
session_key = None
while True:
session_key = self.connect()
if session_key is not None:
break
print("Connection established with", self.DST_ADDR)
#print(session_key)
self.current_protocol = "FMP"
self.session_key = session_key
self.seq_last_sent = 0
while True:
self.fmp_process()