#!/usr/bin/env python2

# Author: Troels Arvin <troels@arvin.dk>
#
# ------------------------------ disclaimer, etc: ------------------------------
# Copyright (C) 2011 by Troels Arvin.
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# ------------------------------ disclaimer, etc: ------------------------------

# This script will flip bits at random positions in a file, which will in most
# cases destroy the file, so use it with caution. The purpose of the script
# is to be able to simulate the effects of random bit errors in a storage 
# system.
#
# If the file is large, then this script should be run in a 64-bit operating
# system. Python version 2.5 or newer is required.
# See http://troels.arvin.dk/code/bit_flipper for the latest version.

import sys
import os
import random
import mmap
from collections import defaultdict

default_num_mutations = 10

this_script = os.path.basename(__file__)
random.seed()

def usage():
    print('Usage:')
    print('  %s <file_name> [number of mutations]' % this_script)
    print('')
    print('Default number of mutations is %d.' % default_num_mutations)
    print('')
    print('This utility will flip a bit at random positions in the file given by')
    print('the <file_name> argument. Use with caution!')

def err(msg):
    sys.stderr.write('Error: %s.%s' % (msg,os.linesep))
    sys.exit(1)

def get_random_numbers(file_size_bits, num_mutations):
    s = set()
    tries = 0
    while len(s) < num_mutations:
        tries += 1
        s.add(random.randint(0,file_size_bits-1))
    return s

def translate_bitnum_to_bytenum_and_bitnum(set_of_bitnums):
    byte_offsets = defaultdict(list)
    for bitnum in set_of_bitnums:
        bytenum = bitnum // 8
        byte_offset = bitnum - (bytenum * 8)
        byte_offsets[bytenum].append(byte_offset)
    return byte_offsets

def flip_in_char(c,offset):
    i = ord(c)
    flipper = 1 << offset
    return chr(i ^ flipper)

num_args = len(sys.argv)
if num_args not in [2,3]:
    usage()
    sys.exit(1)

file_name = sys.argv[1]
if file_name in ['-h','--help']:
    usage()
    sys.exit(0)

num_mutations = default_num_mutations
if num_args == 3:
    try:
        num_mutations = int(sys.argv[2])
    except:
        err('Number of mutations needs to be an integer')

if num_mutations < 1:
    err('Number mutations need to be greater than zero')

file_size = os.path.getsize(file_name)
file_size_bits = file_size * 8
if file_size_bits < num_mutations:
    err('File is too small for %d mutations' % num_mutations)

mutate_bit_positions = get_random_numbers(file_size_bits, num_mutations)
mutate_bytes_and_bits = translate_bitnum_to_bytenum_and_bitnum(mutate_bit_positions)
bytes = mutate_bytes_and_bits.keys()
bytes.sort()

try:
    f = open(file_name,'r+b')
except:
    err('Could not open file %s for reading and writing')

try:
    map = mmap.mmap(f.fileno(),0)
except (Exception, e):
    err('Could not mmap file: %s' % e)

for b in bytes:
    oldchar = map[b]
    newchar = oldchar
    offsets = mutate_bytes_and_bits[b]
    for offset in offsets:
        newchar = flip_in_char(newchar,offset)
    map[b] = (newchar) 
    print('File offset %d flipped at bit-positions %s.' % (b, str(offsets)))

map.close()
f.close()