Refactored server, minor changes in client

Author: skelly

client/src/fraction.c

@@ -104,7 +104,7 @@ void print_fraction(fraction_t fraction) {
   printf("Data size: %lu\n\n", fraction.data_size);
 }
 
-uint32_t calc_crc(fraction_t *frac){
+int calc_crc(fraction_t *frac){
     uint8_t buffer[sizeof(frac->magic) + sizeof(frac->index) + sizeof(frac->iv) + frac->data_size];
     size_t offset = 0;
 
@@ -122,9 +122,7 @@ uint32_t calc_crc(fraction_t *frac){
 
     uint32_t calculated_crc = crc32(buffer, offset);
 
-    if (calculated_crc == frac->crc) {
-        printf("Checksum correct\n");
-    } else {
+    if (calculated_crc != frac->crc) {
         printf("Checksum incorrect\n");
         printf("Checksum generated: %08X\n", calculated_crc);
         printf("Checksum from fraction: %08X\n\n", frac->crc);
@@ -134,14 +132,15 @@ uint32_t calc_crc(fraction_t *frac){
 }
 
 int check_fractions(fraction_t *fraction, size_t size){
+  int res = 0;
   for(size_t i = 0; i < size; i++){
     if (!calc_crc(&fraction[i])) {
       fprintf(stderr, "Failed to validate integrity of fraction:\n");
       print_fraction(fraction[i]);
-      return 1;
+      res += 1;
     }
   }
-  return 0;
+  return res;
 }
 
 void fraction_free(fraction_t *fraction) {

client/src/main.c

@@ -84,11 +84,8 @@ int main(void) {
       fprintf(stderr, "Failed to download fraction\n");
     }
   }
-
+  puts("Downloaded fractions");
   qsort(fractions, lines_read, sizeof(fraction_t), compare_fractions);
-  for (int i = 0; i < lines_read; i++) {
-    print_fraction(fractions[i]);
-  }
 
   if (check_fractions(fractions, lines_read)) { // if this works, s0s4 and skelly is to blame!
     fprintf(stderr, "Fractions check failed\n");
@@ -98,6 +95,7 @@ int main(void) {
     http_free(&http_post_res);
     goto cleanup;
   }
+  puts("Verified fractions");
 
   if (http_post(sfd, "/deadbeef", "plain/text", "{'downloaded':true}",
                 &http_post_res) != HTTP_SUCCESS) {

server/fraction.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass, field
+import zlib
+import struct
+from typing import Literal
+import logging
+
+@dataclass
+class Fraction:
+    """Dataclass to represent a fraction"""
+    magic: int
+    index: int
+    iv: bytes
+    data: bytes
+    
+    _crc: int = field(init=False, repr=False)
+    _generated_crc: bool = field(init=False, repr=False, default=False)
+
+    def __post_init__(self) -> None:
+        """Post-initialization: calculate CRC after dataclass initialization"""
+        self.calculate_crc()
+
+    def header_to_bytes(
+        self, endianess: Literal["big", "little"] = "little", include_crc=True
+    ) -> bytes:
+        """
+        Convert the header information of the fraction to bytes.
+        
+        endianess: Endianness to use (big, little)
+        include_crc: Include CRC in the returned data (default: True)
+        """
+        end = ">" if endianess == "big" else "<"
+        fmt = f"{end}II16s{'I' if include_crc else ''}"
+
+        # Pack the header, optionally including CRC
+        args = [self.magic, self.index, self.iv]
+        if include_crc:
+            args.append(self.crc)
+        
+        header_data = struct.pack(fmt, *args)
+        logging.debug(f"Header data [{self.index}]: {header_data.hex()}")
+        return header_data
+
+    def calculate_crc(self) -> None:
+        """Calculate and store the CRC checksum of the fraction"""
+        crc_data = self.header_to_bytes(include_crc=False) + self.data
+        self._crc = zlib.crc32(crc_data)
+        self._generated_crc = True
+        logging.debug(f"Calculated CRC for fraction [{self.index}]: {self._crc:#x}")
+
+    @property
+    def crc(self) -> int:
+        """Get or calculate the CRC if not already generated"""
+        if not self._generated_crc:
+            self.calculate_crc()
+        return self._crc
+
+    @property
+    def data_size(self) -> int:
+        """Return the size of the data"""
+        return len(self.data)

server/fractionator.py

@@ -1,65 +1,11 @@
 import os
 import io
-from typing import Callable, Optional
+from typing import Optional
 import logging
 import secrets
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 import utils
-from dataclasses import dataclass, field
-import zlib
-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"] = "little", 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 "<"
-        fmt = f"{end}II16sI" if crc else f"{end}II16s"
-    
-        args = [fmt, self.magic, self.index, self.iv]
-        if crc:
-            args.append(self._crc)
-
-        header_data = struct.pack(*args);
-        logging.info(f"Header data: {header_data.hex()}")
-        return header_data
-
-    def calculate_crc(self) -> None:
-        """Calculate the CRC checksum of the fraction"""
-        crc_data = self.header_to_bytes(crc=False) + self.data
-        self._crc = zlib.crc32(crc_data)
-
-    @property
-    def crc(self) -> int:
-        if not self._crc:
-            self.calculate_crc()
-
-        return self._crc
-
-    @property
-    def data_size(self) -> int:
-        return len(self.data)
-    
-    def __post_init__(self) -> None:
-        self.calculate_crc()
-
+from fraction import Fraction
 
 class Fractionator:
     MAGIC: int = 0xDEADBEEF
@@ -82,16 +28,33 @@ def __init__(
         self._key: Optional[bytes] = Fractionator.validate_aes_key(
             key
         )  # AES-256 cryptographic key
+        self._mode = modes.CFB
+        self._algorithm = algorithms.AES256
+        self._cipher = None
+        
+    @property
+    def cipher(self) -> Cipher:
+        if not self._cipher:
+            if not self._key or not self._iv:
+                raise ValueError(f"Missing key or IV (_key:{self._key}, _iv:{self._iv})")
+            
+            self._cipher = Cipher(self._algorithm(self._key), self._mode(self._iv))
+    
+        return self._cipher
 
     def open_reading_stream(self) -> None:
         """
         Opens a reading stream to the file specified in self._path.
         If a stream is already open, this function has no effect
         """
-        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}.")
-            return
+        try:
+            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}.")
+                return
+        except FileNotFoundError as err:
+            logging.error(f"File not found: {self._path}")
+            raise err
 
     def _make_fraction(self, index: int) -> None:
         """Read from the object-file and generate a fraction"""
@@ -103,19 +66,23 @@ def _make_fraction(self, index: int) -> None:
             Fractionator.CHUNK_SIZE
         )  # don't use peek, as it does not advance the position in the file
 
-        # 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
-
+        # generate iv and encrypt the chunk
+        encrypted_data = self._encrypt_chunk(data)
+        
         # Create a fraction instance and add it to self._fractions
         fraction = Fraction(
             magic=Fractionator.MAGIC, index=index, iv=self._iv, data=encrypted_data
         )
         self._fractions.append(fraction)
         logging.debug(f"Created fraction #{fraction.index}")
 
+    def _encrypt_chunk(self, data: bytes) -> bytes:
+        self._iv = self._generate_iv()
+        return self.do_aes_operation(data, True)
+    
+    def _generate_iv(self) -> bytes:
+        return secrets.token_bytes(16)
+    
     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)
@@ -163,11 +130,11 @@ def load_backup(self, backup_path: str):
             with open(backup_path, "r") as f:
                 self.fraction_paths = [line.strip() for line in f]
             logging.debug(
-                f"[debug: _load_backup] Loaded {len(self.fraction_paths)} paths from backup."
+                f"Loaded {len(self.fraction_paths)} paths from backup."
             )
 
         except OSError as e:
-            logging.error(f"[error: _load_backup] Failed to load backup: {e}")
+            logging.error(f"Failed to load backup: {e}")
             return []
 
     def _clean_fraction(self, path: str):
@@ -190,12 +157,13 @@ def clean_fractions(self) -> None:
         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))
+        cipher = self.cipher
         operator = cipher.encryptor() if op else cipher.decryptor()
 
         return operator.update(data) + operator.finalize()

server/utils.py

@@ -1,7 +1,6 @@
 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))