#!/usr/bin/python3

# Disk RW test. Takes one argument, a path for a tmp file
# Checks for file creation/write/seek/read/diff, returns critical on any failure

import argparse
import tempfile
import string
import random

def main():
    data = bytes(''.join(random.choices(string.ascii_uppercase + string.digits, k=20)), "UTF-8")

    fail = []
    for path in args.path:
        # Make tmp file
        try:
            fh = tempfile.TemporaryFile(dir = path)
        except:
            fail.append("can't create temp file in %s" % path)
            continue

        # Write to tmp file
        try:
            fh.write(data)
        except:
            fail.append("can't write file in %s" % path)
            continue

        # Seek tmp file
        try:
            fh.seek(0)
        except:
            fail.append("Can't seek file in %s" % path)
            continue

        # Read tmp file
        try:
            foo = fh.read()
        except:
            fail.append("Can't read file in %s" % path)
            continue

        if foo != data:
            fail.append("File contents doesn't match in %s" % path)
            continue

        # Delete tmp file
        fh.close()

    if (fail):
        print("CRITICAL: %s" % ",".join(fail))
        exit(2)
    else: 
        print("OK: Filesystem(s) writable: %s" % ",".join(args.path))
        exit(0)

def parse_cmdline():
    parser = argparse.ArgumentParser( description="Check if a disk is writable")
    parser.add_argument('path', nargs="*", type=str, default=["/tmp"], help="Path for temp file (default /tmp)")
    args = parser.parse_args()
    return args

args = parse_cmdline()
main()
