diff --git a/src/client/client b/src/client/client new file mode 100755 index 0000000..3b9e52c Binary files /dev/null and b/src/client/client differ diff --git a/src/client/http.c b/src/client/http.c index f5ed4b9..3c704cb 100644 --- a/src/client/http.c +++ b/src/client/http.c @@ -6,12 +6,9 @@ #include #include -#define MIN(X, Y) (((X) < (Y)) ? (X) : (Y)) - -#define HTTP_BUFFER_SIZE 1024 - const char *CONTENT_LENGTH = "Content-Length: "; -const char *GET_REQ_TEMPLATE = "GET %s HTTP/1.1\r\nConnection: keep-alive\r\n\r\n"; +const char *GET_REQ_TEMPLATE = + "GET %s HTTP/1.1\r\nConnection: keep-alive\r\n\r\n"; int http_get(int sfd, const char *path, http_res_t *res) { char buf[HTTP_BUFFER_SIZE]; @@ -22,7 +19,7 @@ int http_get(int sfd, const char *path, http_res_t *res) { buf[HTTP_BUFFER_SIZE - 1] = 0; // ensure buf is null terminated - snprintf(buf, 1023, GET_REQ_TEMPLATE, path); + snprintf(buf, HTTP_BUFFER_SIZE-1, GET_REQ_TEMPLATE, path); send_request(sfd, buf); total_bytes = 0; diff --git a/src/client/http.h b/src/client/http.h index 61bb7b8..3d8a124 100644 --- a/src/client/http.h +++ b/src/client/http.h @@ -3,6 +3,10 @@ #include +#define MIN(X, Y) (((X) < (Y)) ? (X) : (Y)) + +#define HTTP_BUFFER_SIZE 1024 + // error codes #define HTTP_SUCCESS 0 #define HTTP_SOCKET_ERR 1 diff --git a/src/client/main.c b/src/client/main.c index d42f021..dcae651 100644 --- a/src/client/main.c +++ b/src/client/main.c @@ -1,4 +1,3 @@ -#include #include #include #include @@ -28,10 +27,9 @@ int main() { if (HTTP_SUCCESS != http_get(sfd, "/", &fraction_links_resp)) { return EXIT_FAILURE; } - write(1, response.data, response.size); - http_free(&response); - + write(1, fraction_links_resp.data, fraction_links_resp.size); + http_free(&fraction_links_resp); - close(sfd); - return EXIT_SUCCESS; + close(sfd); + return EXIT_SUCCESS; } diff --git a/src/server/fractionator.py b/src/server/fractionator.py index 1eddbba..7cb42ea 100644 --- a/src/server/fractionator.py +++ b/src/server/fractionator.py @@ -10,68 +10,81 @@ import struct from typing import Literal + @dataclass class Fraction: """Dataclass to represent a fraction""" + magic: int index: int iv: bytes _crc: int = field(init=False, repr=False) data: bytes - - def header_to_bytes(self, - endianess: Literal["big", "little"]="big", - crc=True + + def header_to_bytes( + self, endianess: Literal["big", "little"] = "big", crc=True ) -> bytes: """ Convert the header information of the fraction to bytes - + endianess: Endianess to use (big, little) crc: Include CRC in the returned data (default: True) """ - end = ">" if endianess=="big" else "<" + end = ">" if endianess == "big" else "<" fmt = f"{end}II16sI" if crc else f"{end}II16s" - + args = [fmt, self.magic, self.index, self.iv] - if crc: args.append(self._crc) - + if crc: + args.append(self._crc) + return struct.pack(*args) - + def calculate_crc(self) -> None: """Calculate the CRC checksum of the fraction""" crc_data = self.header_to_bytes(crc=False) self._crc = zlib.crc32(crc_data) - + @property def crc(self) -> int: if not self._crc: self.calculate_crc() - + return self._crc - + def __post_init__(self) -> None: - self.calculate_crc() - + self.calculate_crc() + + class Fractionator: - MAGIC: int = 0xdeadbeef + MAGIC: int = 0xDEADBEEF CHUNK_SIZE: int = 8192 FRACTION_PATH_LEN = 16 - - def __init__(self, path: str, out_path: str, key: bytes, backup: str = ".erebos_bckp") -> None: + + def __init__( + self, path: str, out_path: str, key: bytes, backup: str = ".erebos_bckp" + ) -> None: """Class to handle loading/preparation of a Fractionator object file to feed to the loader""" - self._path: str = os.path.abspath(Fractionator.validate_source_path(path)) # Path to Fractionator object file - - self._out_path: str = os.path.abspath(Fractionator.validate_output_path(out_path)) # Path to store generated fractions + self._path: str = os.path.abspath( + Fractionator.validate_source_path(path) + ) # Path to Fractionator object file + + self._out_path: str = os.path.abspath( + Fractionator.validate_output_path(out_path) + ) # Path to store generated fractions self.backup_path = os.path.join(self._out_path, backup) - - self._fractions: list[Fraction] = [] # Keep track of the fraction objects - self._fraction_paths: list[str] = [] # Book-keeping of fraction filenames for cleanup + + self._fractions: list[Fraction] = [] # Keep track of the fraction objects + self._fraction_paths: list[str] = ( + [] + ) # Book-keeping of fraction filenames for cleanup # I/O self._buf_reader: Optional[io.BufferedReader] = None # AES-256 related instance attributes - self._iv: Optional[bytes] = None # AES-256 initialization vector - self._key: Optional[bytes] = Fractionator.validate_aes_key(key) # AES-256 cryptographic key + self._iv: Optional[bytes] = None # AES-256 initialization vector + self._key: Optional[bytes] = Fractionator.validate_aes_key( + key + ) # AES-256 cryptographic key def open_reading_stream(self) -> None: """ @@ -80,40 +93,41 @@ def open_reading_stream(self) -> None: """ if self._buf_reader is None or self._buf_reader.closed: self._buf_reader = open(self._path, "rb") - logging.debug(f"Opened reading stream to {self._path}.") + logging.debug(f"Opened reading stream to {self._path}.") return def _make_fraction(self, index: int) -> None: """Read from the object-file and generate a fraction""" - if not isinstance(index, int): + if not isinstance(index, int): raise ValueError(f"index must be an integer (got `{type(index)}`)") # Open a stream to the file and read a chunk self.open_reading_stream() - data = self._buf_reader.read(Fractionator.CHUNK_SIZE) # don't use peek, as it does not advance the position in the file + data = self._buf_reader.read( + Fractionator.CHUNK_SIZE + ) # don't use peek, as it does not advance the position in the file # logging.debug("[debug: _make_fraction] Read chunk from stream.") - + # Generate an IV and encrypt the chunk - self._iv = secrets.token_bytes(16) # initialization vector for AES-256 encryption - encrypted_data = self.do_aes_operation(data, True) # encrypt chunk + self._iv = secrets.token_bytes( + 16 + ) # initialization vector for AES-256 encryption + encrypted_data = self.do_aes_operation(data, True) # encrypt chunk # logging.info("[info: _make_fraction] Encrypted chunk data using AES-256") - + # Create a fraction instance and add it to self._fractions fraction = Fraction( - magic=Fractionator.MAGIC, - index=index, - iv=self._iv, - data=encrypted_data + magic=Fractionator.MAGIC, index=index, iv=self._iv, data=encrypted_data ) self._fractions.append(fraction) # logging.debug(f"[debug: _make_fraction] Created Fraction object: {fraction} (crc: {fraction.crc})") logging.debug(f"Created fraction #{fraction.index}") - + def make_fractions(self) -> None: """Iterate through the Fractionator object file specified in self._path and generate Fraction objects""" size = os.path.getsize(self._path) num_chunks = (size + Fractionator.CHUNK_SIZE - 1) // Fractionator.CHUNK_SIZE - + logging.info(f"[info: make_fractions] Creating {num_chunks} fractions.") for i in range(num_chunks): self._make_fraction(i) @@ -121,23 +135,25 @@ def make_fractions(self) -> None: def _write_fraction(self, fraction: Fraction): """Write a fraction to a file""" os.makedirs(self._out_path, exist_ok=True) - path = os.path.join(self._out_path, utils.random_string(Fractionator.FRACTION_PATH_LEN)) - + path = os.path.join( + self._out_path, utils.random_string(Fractionator.FRACTION_PATH_LEN) + ) + with open(path, "wb") as f: header_data = fraction.header_to_bytes() data = fraction.data - + f.write(header_data) f.write(data) - + self._fraction_paths.append(path) logging.debug(f"Wrote fraction #{fraction.index} to {path}") - + def write_fractions(self) -> None: """Convert the fraction objects to pure bytes and write them in the appropriate directory (self._out)""" for fraction in self._fractions: self._write_fraction(fraction) - + if self.backup_path: self._save_backup() @@ -150,14 +166,15 @@ def _save_backup(self) -> None: logging.debug(f"Backup saved at {self.backup_path}.") except OSError as e: logging.error(f"Failed to save backup: {e}") - - + def _load_backup(self) -> list[str]: """Load fraction paths from the backup file.""" try: with open(self.backup_path, "r") as f: paths = [line.strip() for line in f] - logging.debug(f"[debug: _load_backup] Loaded {len(paths)} paths from backup.") + logging.debug( + f"[debug: _load_backup] Loaded {len(paths)} paths from backup." + ) return paths except OSError as e: logging.error(f"[error: _load_backup] Failed to load backup: {e}") @@ -170,25 +187,25 @@ def _clean_fraction(self, path: str): logging.debug(f"Removed {path}.") except FileNotFoundError: logging.debug(f"File not found: {path}") - + def clean_fractions(self) -> None: logging.info("Cleaning fractions . . .") if self.backup_path and not self._fraction_paths: self._fraction_paths = self._load_backup() - + if not self._fraction_paths: logging.error("No fraction paths detected.") for path in self._fraction_paths: self._clean_fraction(path) - + self._fraction_paths = [] logging.info("Done.") - + def do_aes_operation(self, data: bytes, op: bool) -> bytes: """Perform an AES-256 operation on given data (encryption [op=True]/decryption [op=False])""" if not self._key or not self._iv: raise ValueError(f"Missing key or IV (_key:{self._key}, _iv:{self._iv})") - + cipher = Cipher(algorithms.AES(self._key), modes.OFB(self._iv)) operator = cipher.encryptor() if op else cipher.decryptor() @@ -201,23 +218,24 @@ def _close_stream(self) -> None: self._buf_reader = None logging.debug(f"Closed stream to {self._path}.") return - + logging.debug(f"No stream was open.") @staticmethod def validate_aes_key(key: bytes) -> bytes: """Check if key is a valid AES-256 key (32 bytes)""" if not isinstance(key, bytes) or len(key) != 32: - raise ValueError(f"Invalid AES-256 key. (expected 32 bytes of `{bytes}`, got {len(key)} of `{type(key)}`)") + raise ValueError( + f"Invalid AES-256 key. (expected 32 bytes of `{bytes}`, got {len(key)} of `{type(key)}`)" + ) return key - - + @staticmethod def validate_file_ext(path: str, extension: str) -> str: """Checks if path is a file and ends with extension""" if not path.endswith(".ko") or not os.path.isfile(path): raise ValueError(f"{path} is not a valid file.") - + return path @staticmethod @@ -226,17 +244,17 @@ def validate_source_path(path: str) -> str: if not os.path.exists(path): raise FileNotFoundError("Path not found.") path = Fractionator.validate_file_ext(path, ".ko") - + return path - + @staticmethod def validate_output_path(path: str) -> str: """Checks if path exists and is a directory. it will create a new directory otherwise""" os.makedirs(path, exist_ok=True) if not os.path.isdir(path): raise ValueError(f"Path is not a directory ({path}).") - + return path - + def __del__(self) -> None: - self._close_stream() \ No newline at end of file + self._close_stream() diff --git a/src/server/main.py b/src/server/main.py index 3b2b822..2c7b794 100644 --- a/src/server/main.py +++ b/src/server/main.py @@ -16,6 +16,7 @@ BACKUP_FILENAME = ".erebos_bckp" + def handle_args(parser: argparse.ArgumentParser): """Configure the given ArgumentParser""" parser.add_argument( @@ -47,6 +48,7 @@ def handle_args(parser: argparse.ArgumentParser): "--rm-backup", action="store_true", help="Remove the generated backup file" ) + if __name__ == "__main__": parser = argparse.ArgumentParser() handle_args(parser) @@ -72,4 +74,4 @@ def handle_args(parser: argparse.ArgumentParser): lkm.write_fractions() # Stage fractions over HTTP - start_server(args.bind, args.port) \ No newline at end of file + start_server(args.bind, args.port) diff --git a/src/server/server.py b/src/server/server.py index ba1f748..09ec5d4 100644 --- a/src/server/server.py +++ b/src/server/server.py @@ -41,11 +41,13 @@ def list_directory(self, path): file_list.sort(key=lambda a: a.lower()) contents = [] - + enc = sys.getfilesystemencoding() + server_addr = self.server.server_address + host, port = server_addr for name in file_list: - display_name = f"http://{self.headers['Host']}{self.path}{name}" + display_name = f"http://{host}:{port}{self.path}{name}" contents.append(html.escape(display_name, quote=False)) encoded = "\n".join(contents).encode(enc, "surrogateescape") @@ -57,12 +59,13 @@ def list_directory(self, path): self.send_header("Content-Length", str(len(encoded))) self.end_headers() return f - + + def start_server(bind, port): test( HandlerClass=PlainListingHTTPRequestHandler, ServerClass=DualStackServer, - protocol="HTTP/1.1", # permit keep-alive connections + protocol="HTTP/1.1", # permit keep-alive connections port=port, bind=bind, ) diff --git a/src/server/utils.py b/src/server/utils.py index b5f39ea..7bd6041 100644 --- a/src/server/utils.py +++ b/src/server/utils.py @@ -1,6 +1,7 @@ import random import string + def random_string(n: int = 16, sample: str = string.ascii_lowercase + string.digits): """Returns a random string using the characters defined in sample""" - return "".join(random.choices(sample, k=n)) \ No newline at end of file + return "".join(random.choices(sample, k=n))