Simple utility to create a(n optionally encrypted) ramdisk on macos catalina. Python 2 & 3 compatible

ramdisk.py

#!/usr/bin/env python

# Copyright 2020 Jonathan Simms
#
# Redistribution and use in source and binary forms, with or without modification, are permitted
# provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this list of conditions
# and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice, this list of
# conditions and the following disclaimer in the documentation and/or other materials provided with
# the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors may be used to
# endorse or promote products derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
# IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from __future__ import print_function

import argparse
import os
import sys

from subprocess import check_output, check_call
from tempfile import TemporaryFile

PASSPHRASE_SIZE=127

HELP = """\
if '--crypted' is given, an encrypted APFS volume with a (max length) 127 byte
random passphrase will be created. Note you will need to eject 2 devices the
apfs container *and* the ramdisk to clean up.

if 'volname' is given it will be used as the volume name, otherwise it will be
called 'RAMDisk'

"""

def main():
  args = sys.argv[1:]
  if len(args) < 1:
    print(HELP, file=sys.stderr)
    sys.exit(1)

  ap = argparse.ArgumentParser(
      formatter_class=argparse.RawDescriptionHelpFormatter,
      description=HELP)
  ap.add_argument('size_in_mb', type=int, help="the size of the ramdisk to create in mb")
  ap.add_argument(
    '-V', '--volname',
    type=str,
    help="the name of the volume to create",
    default="RAMDisk")
  ap.add_argument(
    '--crypted',
    action='store_true',
    help="create a crypted APFS volume with a random passphrase (127B read from /dev/urandom)")

  opts = ap.parse_args()

  assert opts.size_in_mb > 0

  def mk_crypted(base_dev):
    try:
      output = check_output(["diskutil", "apfs", "createContainer", base_dev]).decode('utf8')
      info = [line for line in output.split("\n") if line.startswith('Disk from APFS operation: ')]
      if len(info) < 1:
        raise RuntimeError("could not find created container device name")
      dev2 = info[0].split(': ', 2)[1].rstrip()

      with TemporaryFile() as tmp:
        random = os.urandom(PASSPHRASE_SIZE)
        tmp.write(random)
        tmp.flush()
        os.fsync(tmp.fileno())
        tmp.seek(0, os.SEEK_SET)
        assert tmp.tell() == 0
        check_call(
          ['diskutil', 'apfs', 'addVolume', dev2, 'APFS', opts.volname, '-stdinpass'],
          stdin=tmp
        )
      return dev2
    except Exception as e:
      check_call(['diskutil', 'eject', base_dev])
      raise e

  devices = []

  dev1 = check_output(
      ['hdiutil', 'attach', '-nomount', 'ram://' + str(2048 * opts.size_in_mb)]
    ).decode("utf8") \
     .rstrip()

  devices.append(dev1)

  if opts.crypted:
    devices.append(mk_crypted(dev1))
  else:
    check_call(['diskutil', 'erasevolume', 'HFS+', opts.volname.strip(), dev1])

  cmd = ' && '.join(["diskutil eject %s" % (d,) for d in reversed(devices)])

  print("to clean up run command: '%s'" % (cmd,))


if __name__ == '__main__':
  main()