Detect fake USB drives

fake-usb-tester.py

#!/usr/bin/env python3
#
# WARNING: This program will OVERWRITE the entire device and
#          you may LOSE DATA. Use it carefully.

# How to create a test block device:
#   $ dd if=/dev/zero of=fake bs=1M count=100
#   # mknod fake b 7 200
#   # losetup /dev/fake fake

import tqdm

import argparse
import os
import random

parser = argparse.ArgumentParser(description="Detect fake USB drives")
parser.add_argument("device",
                    help="Path to the device (like /dev/fake)")
parser.add_argument("size", default=8, type=int, nargs='?',
                    help="Block size (in MB)")
args = parser.parse_args()
block_size = args.size*1024**2

def new_chunk(size):
    return bytearray(random.getrandbits(8) for _ in range(size))

# Get device size
fd = os.open(args.device, os.O_RDONLY)
size = os.lseek(fd, 0, os.SEEK_END)
os.close(fd)

def chunk_iterator(**kwargs):
    return tqdm.trange(int(size/block_size),
                       unit='block',
                       **kwargs)

state = random.getstate()
with open(args.device, 'wb') as device:
    for test in chunk_iterator(desc='Writing'):
        device.write(new_chunk(block_size))

random.setstate(state)
with open(args.device, 'rb') as device:
    for test in chunk_iterator(desc='Reading'):
        if new_chunk(block_size) != device.read(block_size):
            print("Fake device found!")
            break
    else:
        print("Device seems legit.")