Enable/Disable ASLR on Linux

aslr.py

#!/usr/bin/env python
"""
mv aslr.py /usr/bin/aslr && chmod +x /usr/bin/aslr

Be careful!
"""
import os, sys

FLAG_ENABLE = 2
FLAG_DISABLE = 0
FILENAME = '/proc/sys/kernel/randomize_va_space'

def _write_flag(filename, flag):
    with open(filename, 'wb') as fd:
        fd.write("%s" % (str(flag)))


def enable():
    _write_flag(FILENAME, FLAG_ENABLE)

def disable():
    _write_flag(FILENAME, FLAG_DISABLE)

def main(option):
    options = {
        'on': enable,
        'enable': enable,
        'off': disable,
        'disable': disable
        }

    try:
        func = options[option]
    except KeyError:
        print("Invalid option! (%s)" % ('|'.join(options.keys())))
        quit(1)

    func()


if __name__ == '__main__':
    if os.geteuid() is not 0:
        print("You must call this program as root!")
        quit(1)
        
    try:
        opt = sys.argv[1]
    except IndexError:
        opt = None

    main(opt)