ddnuke: add fdisk+prompt to startup, and rough ETA

[home.git] / ddnuke.c

/*
 * released into the public domain
 * written by Mike Frysinger <vapier>
 */

#pragma GCC diagnostic warning "-Wall"
#pragma GCC diagnostic warning "-Wextra"

#define _FILE_OFFSET_BITS 64
#define _LARGEFILE_SOURCE
#define _LARGEFILE64_SOURCE
#define _GNU_SOURCE

#include <ctype.h>
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <locale.h>
#include <spawn.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sysexits.h>
#include <time.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <linux/fs.h>

static unsigned long long mbps(struct timespec *stime, struct timespec *etime, off_t len)
{
	uint64_t dtime;
	clock_gettime(CLOCK_MONOTONIC, etime);
	dtime = ((etime->tv_sec - stime->tv_sec) * 1000 * 1000 * 1000) +
		(etime->tv_nsec - stime->tv_nsec);
	*stime = *etime;
	return 1000 * len / dtime;
}

static off_t get_blk_size(int fd)
{
	uint64_t size;
	if (ioctl(fd, BLKGETSIZE64, &size))
		err(EX_DATAERR, "ioctl(BLKGETSIZE64) failed");
	return size;
}

static off_t get_size(int fd)
{
	struct stat st;

	if (fstat(fd, &st))
		err(EX_DATAERR, "could not stat %i", fd);

	if (S_ISREG(st.st_mode))
		return st.st_size;
	else if (S_ISBLK(st.st_mode))
		return get_blk_size(fd);

	errno = EINVAL;
	err(EX_DATAERR, "unknown type of file");
}

static void nuke(int fd, off_t offset, off_t max_size, unsigned char pattern)
{
	static char mem[1024 * 1024];
	memset(mem, pattern, sizeof(mem));
	if (pattern)
		printf("Writing 0x%X to the output\n", pattern);

	if (lseek(fd, offset, SEEK_SET) != offset)
		err(EX_DATAERR, "lseek(%"PRIu64"u) failed", offset);

	unsigned long long speed = 0;
	struct timespec stime, etime, itime;
	clock_gettime(CLOCK_MONOTONIC, &stime);
	itime = stime;

	off_t pos = offset, last_pos = 0;
	ssize_t ret;
	int fsync_pos = 0;
	while (pos < max_size) {
		/* This will round up to sizeof(mem) ... */
		ret = write(fd, mem, sizeof(mem));
		pos += ret;
		if (ret != sizeof(mem)) {
			if (pos != max_size)
				printf("\nread() returned %zi (wanted %zu)\n%'llu MB/s over entire run\n",
					ret, sizeof(mem), mbps(&itime, &etime, pos));
			return;
		}

		unsigned long long eta = speed ? (max_size - pos) / (speed * 1000 * 1000) : 0;
		printf("\r\e[2K%'llu bytes %u%% (%'llu MB/s) (ETA ~%llum%llus)",
			(unsigned long long)pos, (unsigned)((pos * 100) / max_size), speed,
			eta / 60, eta % 60);

		if ((++fsync_pos % 32) == 0) {
			speed = mbps(&stime, &etime, pos - last_pos);
			last_pos = pos;
			fflush(stdout);
			fsync(fd);
		}
	}

	printf("\n");
}

static void usage(int status)
{
	fprintf(
		status ? stderr : stdout,
		"Usage: ddnuke [options] <dev>\n"
		"\n"
		"Options:\n"
		" -q            Run quietly\n"
		" -y            Assume yes to all prompts\n"
		" -o <offset>   Position to start writing\n"
		" -r            Write three times: 0x00, then 0xaa, then 0x55\n"
	);
	exit(status);
}

int main(int argc, char *argv[])
{
	off_t offset = 0;
	const char *file;
	bool quiet = false;
	bool random = false;
	bool yes = false;
	int o;

	setlocale(LC_NUMERIC, "en_US");

	while ((o = getopt(argc, argv, "ho:qry")) != -1) {
		switch (o) {
		case 'o':
			offset = atoll(optarg);
			break;
		case 'q':
			quiet = true;
			break;
		case 'r':
			random = true;
			break;
		case 'y':
			yes = true;
			break;
		case 'h':
			usage(EX_OK);
			break;
		default:
			usage(EX_USAGE);
			break;
		}
	}
	if (argc != optind + 1)
		usage(EX_USAGE);
	file = argv[optind];

	int fd = open(file, O_WRONLY|O_CLOEXEC);
	if (fd == -1)
		err(EX_NOINPUT, "open(%s) failed", file);

	if (!quiet) {
		pid_t pid;
		const char *argv[] = {"fdisk", "-l", file, NULL};
		if (posix_spawnp(&pid, argv[0], NULL, NULL, (char *const *)argv, NULL))
			err(EX_OSERR, "spawn(%s %s) failed", argv[0], file);
		waitpid(pid, NULL, 0);
		printf("\n");
	}

	if (!yes) {
		printf("Do you want to wipe '%s'? (y/N) ", file);
		fflush(stdout);
		char c = 0;
		if (fread(&c, 1, 1, stdin)) {/* don't care */}
		if (tolower(c) != 'y')
			errx(EX_UNAVAILABLE, "User decided not to wipe");
	}

	off_t max_size = get_size(fd);

	nuke(fd, offset, max_size, 0x00);
	if (random) {
		nuke(fd, offset, max_size, 0xaa);
		nuke(fd, offset, max_size, 0x55);
	}

	return EX_OK;
}