/
addSteg.py
115 lines (91 loc) · 3.34 KB
/
addSteg.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
'''
Adds steganography to given dataset
'''
import argparse
from stegano import lsb, lsbset
from stegano.lsbset import generators
from tqdm import tqdm
from essential_generators import DocumentGenerator
import os
import multiprocessing
from PIL import Image
import sys
def get_args():
'''
Gets args for this program
param:
None
return:
parsed arguments object
'''
parser = argparse.ArgumentParser()
parser.add_argument('-data', default='data/',
help='Relative data path to main data folder', type=str)
parser.add_argument(
'-threads', help='Number of threads to use for parallel processing', type=int)
parser.add_argument('-bpp', default=1.0,
help='Bits of hidden data per pixel', type=float)
parser.add_argument('-generators', nargs='+',
help='Names of generators to use')
return(parser.parse_args())
# source: https://stackoverflow.com/a/37233621
def walk_dir(input_path):
'''
Walks through subdirectories to grab all necessary data
param:
input_path - str path to parent folder containing all data
return:
path - str path to file containing data
'''
for subdir, dirs, files in os.walk(input_path):
for name in files:
if('.png' in name and 'steg' not in name):
yield os.path.abspath(os.path.join(subdir, name))
def add_steg(path, bpp):
'''
Add steg to each image in directory
param:
path - str path to data folder containing .pngs
bpp - float desired bits per pixel
return:
None
'''
filecounter = len([path for path in walk_dir(path)])
gen = DocumentGenerator()
if(args.generators):
lsb_generators = [getattr(generators, entry)
for entry in args.generators]
count = 0
for filepath in tqdm(walk_dir(path), total=filecounter, unit='files'):
try:
im = Image.open(filepath)
im_size = im.size[0] * im.size[1]
# average sentence length is 60 bytes -> 480 bits
message = '\n'.join([gen.sentence()
for i in range(int(im_size * bpp / 480))])
# remove non-ascii characters as they mess up steganography
message = message.encode('ascii', 'ignore').decode()
if(args.generators):
# + 1 to add normal lsb
gen_index = int(count % (len(lsb_generators) + 1))
if(gen_index == len(lsb_generators)):
secret = lsb.hide(filepath, message)
secret.save(filepath.replace('.png', '.steg.png'))
else:
lsb_gen = lsb_generators[gen_index]
secret = lsbset.hide(filepath, message, lsb_gen())
secret.save(filepath.replace(
'.png', f".{args.generators[gen_index]}.png"))
else:
secret = lsb.hide(filepath, message)
secret.save(filepath.replace('.png', '.steg.png'))
count += 1
except Exception as e:
print(e)
print(filepath)
if __name__ == '__main__':
args = get_args()
for i in range(1, args.threads+1):
p = multiprocessing.Process(target=add_steg, args=(
args.data + 'dir_' + str(i), args.bpp))
p.start()