From ba7c493208d8839eb547336cef21cea50bf62690 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Valentin=20Mogu=C3=A9rou?= <valentin@kaz.bzh>
Date: Sat, 19 Mar 2022 20:45:17 +0100
Subject: [PATCH] Add insert and extract commands

In order to be able to put and read binary files in an image
---
 libstegano.py |  1 +
 steganocli.py | 40 +++++++++++++++++++++++++++++++++++++++-
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/libstegano.py b/libstegano.py
index a6efbcb..06b94b3 100644
--- a/libstegano.py
+++ b/libstegano.py
@@ -63,6 +63,7 @@ def decode_string(bytelist: bytearray) -> str:
 def write_band(band: bytearray, byte_list: bytearray, starting_pos=0) -> bool:
     "Writes a byte sequence in the given band"
 
+    index = starting_pos
     for index in range(starting_pos, len(band), 8):
         band[index:index+8] = (write_lsb_in_bin(band[index+i], bit) for i, bit in enumerate(split_byte_into_bits(byte_list[index//8])))
         if byte_list[index//8] == 0:
diff --git a/steganocli.py b/steganocli.py
index eb87ca5..671050b 100644
--- a/steganocli.py
+++ b/steganocli.py
@@ -50,6 +50,33 @@ def write(args):
 
         write_file(from_file, to_file, byte_list, verbose=args.verbose)
 
+def extract(args):
+    if args.verbose:
+        print("Extract mode enabled.")
+    
+    with open(args.from_file, 'rb') as from_file, open(args.to_file, 'wb') as to_file:
+        byte_list = read_file(from_file, verbose=args.verbose)
+
+        if args.verbose:
+            print(f"Read {(len(byte_list)+1)*8} bits, {len(byte_list)+1} bytes.")
+            print("================== BYTE STREAM ==================")
+            hex_print(byte_list, margin=1)
+            print("================== BYTE STREAM ==================")
+        
+        to_file.write(byte_list)
+
+def insert(args):
+    with open(args.from_file, 'rb') as from_file, open(args.content_file, 'rb') as content_file, open(args.to_file, 'wb') as to_file:
+        byte_list = content_file.read()+b'\0'
+
+        if args.verbose:
+            print("Insert mode enabled.")
+            print("================== BYTE STREAM ==================")
+            hex_print(byte_list, margin=1)
+            print("================== BYTE STREAM ==================")
+
+        write_file(from_file, to_file, byte_list, verbose=args.verbose)
+
 def main():
     parser = ArgParser()
     parser.add_argument('-v', '--verbose', action='store_true')
@@ -60,12 +87,23 @@ def main():
     parser_read.add_argument('from_file', type=Path)
     parser_read.set_defaults(func=read)
 
-    parser_write = subparsers.add_parser('write', help='Dissimulate data in an image')
+    parser_write = subparsers.add_parser('write', help='Dissimulate a string in an image')
     parser_write.add_argument('string', type=str)
     parser_write.add_argument('to_file', type=Path)
     parser_write.add_argument('--from', type=Path, dest='from_file', default=Path('default_image.jpg'))
     parser_write.set_defaults(func=write)
 
+    parser_read = subparsers.add_parser('extract', help='Retrieve data from an image')
+    parser_read.add_argument('from_file', type=Path)
+    parser_read.add_argument('to_file', type=Path)
+    parser_read.set_defaults(func=extract)
+
+    parser_write = subparsers.add_parser('insert', help='Dissimulate data in an image')
+    parser_write.add_argument('content_file', type=Path)
+    parser_write.add_argument('to_file', type=Path)
+    parser_write.add_argument('--from', type=Path, dest='from_file', default=Path('default_image.jpg'))
+    parser_write.set_defaults(func=insert)
+
     args = parser.parse_args()
     args.func(args)