def create_update(update_path, major, minor):
version = updater.Version(major, minor)
try:
update_filepath = settings.UPDATE_PATH
update_filepath = f"{update_filepath}.{version}"
update_file = open(update_filepath, "wb")
except PermissionError:
print(
f"Failed to write update file {settings.UPDATE_PATH}. Insufficient permission or file is locked."
)
return False
# Creating update zip
with update_file:
update_zip = zipfile.ZipFile(update_file, "w")
zip_directory(update_path, update_zip)
update_zip.close()
# Updating registry with update info
version_registry = version.get_update_registry_path()
registry.set_value(version_registry, os.path.abspath(update_filepath))
version.update_current_version()
print(f"Created update version {version} successfully!")
return True
def generate_rsa_keys(display_keys=True):
key_pair = RSA.generate(bits=settings.RSA_KEY_SIZE)
# Adds the keys to settings file
rsa_settings = dict(
RSA_MODULO=hex(key_pair.n),
PUBLIC_KEY=hex(key_pair.e),
)
if not add_to_settings(rsa_settings):
print(f"Failed to add rsa key to {DEFAULT_SETTINGS_PATH}")
return False
# Adds the keys to registry
if not registry.exists(settings.REGISTRY_PATH):
registry.create_key(settings.REGISTRY_PATH)
registry.set_value(settings.RSA_MODULO_REGISTRY, hex(key_pair.n))
registry.set_value(settings.RSA_PUBLIC_REGISTRY, hex(key_pair.e))
registry.set_value(settings.RSA_PRIVATE_REGISTRY, hex(key_pair.d))
# Displays the keys to the user
print("RSA keys were generated!")
print("")
if display_keys:
show_rsa_keys()
return True
def handle_server_update(self):
if len(self.message) != messages.SERVER_UPDATE_MESSAGE.sizeof():
# The message has an incorrect size...
logging.warning(f"Incorrect message size: {len(self.message)}")
return
# Parse the message
try:
message = messages.SERVER_UPDATE_MESSAGE.parse(self.message)
except construct.ConstructError:
# Should never occur
logging.critical(
f"Failed to parse server update message: {self.message.hex()}",
exc_info=True)
return
# Check the running id is more updated than the current id
current_id = registry.get_value(settings.ADDRESS_ID_REGISTRY)
if current_id >= message.address_id:
# An outdated message, should ignore...
logging.info(
f"Received an outdated server update message. current id: {current_id}, received id: {message.address_id}."
)
return
# Checks if the sender requested to spread this message using broadcast
if message.spread:
self.broadcast_message()
# The message is updated, update our data!
address = message.address.decode("ascii")
registry.set_value(settings.UPDATING_SERVER_REGISTRY, address)
registry.set_value(settings.PORT_REGISTRY, message.port)
registry.set_value(settings.ADDRESS_ID_REGISTRY, message.address_id)
logging.info(
f"Updated address to {address} and port to {message.port}")
# Since port could have changed, we restart our socket
self.cleanup_listener()
self.setup_listener()
def init_registry():
if not registry.exists(settings.REGISTRY_PATH):
registry.create_key(settings.REGISTRY_PATH)
if not registry.exists(settings.AUTO_INSTALLATIONS_REGISTRY):
registry.set_value(settings.AUTO_INSTALLATIONS_REGISTRY,
settings.AUTO_INSTALLATIONS)
if not registry.exists(settings.UPDATING_SERVER_REGISTRY):
registry.set_value(settings.UPDATING_SERVER_REGISTRY,
settings.UPDATING_SERVER)
if not registry.exists(settings.PORT_REGISTRY):
registry.set_value(settings.PORT_REGISTRY, settings.PORT)
if not registry.exists(settings.RSA_MODULO_REGISTRY):
registry.set_value(settings.RSA_MODULO_REGISTRY,
settings.RSA_MODULO)
if not registry.exists(settings.RSA_PUBLIC_REGISTRY):
registry.set_value(settings.RSA_PUBLIC_REGISTRY,
settings.PUBLIC_KEY)
if not registry.exists(settings.UPDATE_MAJOR_REGISTRY):
registry.set_value(settings.UPDATE_MAJOR_REGISTRY,
settings.UPDATE_MAJOR)
if not registry.exists(settings.UPDATE_MINOR_REGISTRY):
registry.set_value(settings.UPDATE_MINOR_REGISTRY,
settings.UPDATE_MINOR)
if not registry.exists(settings.VERSION_MAJOR_REGISTRY):
registry.set_value(settings.VERSION_MAJOR_REGISTRY,
settings.VERSION_MAJOR)
if not registry.exists(settings.VERSION_MINOR_REGISTRY):
registry.set_value(settings.VERSION_MINOR_REGISTRY,
settings.VERSION_MINOR)
if not registry.exists(settings.ADDRESS_ID_REGISTRY):
registry.set_value(settings.ADDRESS_ID_REGISTRY,
settings.ADDRESS_ID)
if not registry.exists(settings.SETTINGS_REGISTRY):
registry.set_value(settings.SETTINGS_REGISTRY,
settings.SETTINGS_PATH)
def update_installed_version(self):
registry.set_value(settings.VERSION_MAJOR_REGISTRY, self.major)
registry.set_value(settings.VERSION_MINOR_REGISTRY, self.minor)
def update_current_version(self):
registry.set_value(settings.UPDATE_MAJOR_REGISTRY, self.major)
registry.set_value(settings.UPDATE_MINOR_REGISTRY, self.minor)
def download_update(self, message):
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
# Creates the TCP server that receives the update
listener.settimeout(settings.CONNECTION_TIMEOUT)
listener.bind(("0.0.0.0", 0)) # Bind to random port
port = listener.getsockname()[1]
listener.listen(1)
# Build the request version message
requested_version = Version(message.major, message.minor)
request_version_dict = dict(type=MessageType.REQUEST_VERSION,
crc32=0,
listening_port=port,
major=requested_version.major,
minor=requested_version.minor)
try:
request_version_message = messages.REQUEST_VERSION_MESSAGE.build(
request_version_dict)
# Update CRC32
request_version_dict["crc32"] = messages.calculate_crc(
request_version_message)
request_version_message = messages.REQUEST_VERSION_MESSAGE.build(
request_version_dict)
except construct.ConstructError:
# Should never occur
logging.critical(f"Failed to build request update message",
exc_info=True)
return False
# Creating the file for the update
update_filepath = settings.UPDATE_PATH
update_filepath = f"{update_filepath}.{requested_version}"
update_file = open(update_filepath, "wb")
# Request the update (so that the Updater will connect to our listening socket)
port = registry.get_value(settings.PORT_REGISTRY)
temp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
temp_socket.sendto(request_version_message, 0,
(self.sender[0], port))
temp_socket.shutdown(2)
temp_socket.close()
# Awaits for sender to connect
receiver, _ = listener.accept()
receiver.settimeout(settings.CONNECTION_TIMEOUT)
data_received = 0
hash_object = settings.HASH_MODULE()
# Download the update
logging.info(f"Downloading update of version {requested_version}")
with update_file:
while data_received < message.size:
chunk = receiver.recv(settings.VERSION_CHUNK_SIZE)
if len(
chunk
) == 0: # Occurs when the other side closed the conenction
break
update_file.write(chunk)
hash_object.update(chunk)
data_received += len(chunk)
# Close the TCP connection
receiver.shutdown(socket.SHUT_RD)
receiver.close()
# Validate the update signature
if not rsa_signing.validate_hash(hash_object,
message.update_signature):
# Delete this invalid update file
os.remove(update_filepath)
logging.info(
"Invalid signature for update file (maybe tampered?)")
return False
# Update the registry with the current update
version_registry = requested_version.get_update_registry_path()
registry.set_value(version_registry,
os.path.abspath(update_filepath))
requested_version.update_current_version()
logging.info(f"Received new update: version {requested_version}")
except socket.timeout:
# Connection was timed-out, too bad... abort
logging.info("Connection timed out")
return False
except socket.error:
# Socket error
logging.info("Socket error has occurred")
return False
finally:
listener.close()
return True
def update_auto_installations(values):
if values["-AUTO-"] is True:
registry.set_value(settings.AUTO_INSTALLATIONS_REGISTRY, 1)
else:
registry.set_value(settings.AUTO_INSTALLATIONS_REGISTRY, 0)