Wondering about your COVID-19 risk level and seeking precautionary steps?

Check Now
  • Rally
  • Having Fun Scripting with Ammonite

Having Fun Scripting with Ammonite

By Mark "Justin" Waks | September 15, 2020

Introduction, aka The Problem to Solve

The product I work on (Rally Recover) has a lot of features, one of the primary ones being messaging -- we make it easy for patients to communicate with their healthcare providers. It was adapted from a previous product, which used Layer as the underlying chat system.

In October 2019, we were faced with a problem: Layer was going out of business. We had already shifted to a different messaging service, but we had a lot of existing data living in Layer, which we needed to preserve. I was tasked with fetching a full historical archive of that data, preferably with a few checkpoints, as close to the shutdown date as possible.

In principle, this was straightforward -- Layer had tools to do such downloads. In practice, though, it was a persnickety process, involving several REST calls over the course of several hours. (Not to mention the problem of decoding the tightly-encrypted data.) After manually figuring out how to do it, I realized that trying to repeat that process by hand every day was going to make me crazy: it was complicated and error-prone. Time for a script!

That said -- I'm a Scala guy, not a shell scripter. Scala is deep in my habits and instincts, and I just plain like it a lot, knowing that the strong types save me from many mistakes. I knew that it was possible to write scripts in Scala using Ammonite; this seemed like the time to try it out.

This article is aimed at moderately experienced Scala programmers who occasionally need to do some scripting. It won't teach any of the Scala language, but will show how Ammonite and related tools can make your life a lot easier.

I'll include many code snippets below, but you can find the full script at place location of the script here. This particular script isn't useful any more (since Layer is now defunct), but I encourage you to use it as inspiration for similar problems.

What is Ammonite?

Ammonite is a collection of closely-related libraries and tools, written by Li Haoyi (author of many Scala tools besides these), that help you do "shell-like" things in Scala. It includes:

In this example, we're going to make use of the second and third of those, as well as show off a couple more of the libraries that you can pull in to do powerful stuff very easily. It may be useful to have the full script open next to this article, so you can see where the code snippets I'm talking about fit in context.

Note that, while you will see some new functions here, Ammonite is just Plain Old Scala. It is enhanced with a bunch of tooling to make scripting easier, but all of Scala's functionality works here as you would expect.

Installing and Running Ammonite

If you are on a Mac and have Homebrew installed, you should be able to run:

$ brew install ammonite-repl

Otherwise, check the documentation for the best way to install it.

Either way, that will install the amm command locally. You can simply invoke that command to get a pretty, interactive REPL (with syntax coloring and generally friendlier than the standard Scala REPL), but we're going to focus on writing scripts with it.

To create a script, create a file with the .sc extension (instead of the usual .scala) and start putting code into it. You don't need to create a project or involve sbt or anything like that: you just write Scala code and run it immediately. That's part of the delight of Ammonite scripting: there's very little ceremony.

Given a file like our layer.sc example, we would run it like this:

amm ./layer.sc

But let's make this easier. Like most scripting systems, Ammonite works with the "shebang" mechanism that most modern shells have. Therefore, we can say as the first line of the script:

#!/usr/bin/env amm

Now, at the command line we can say:

layer.sc

The shell will see the shebang, interpret that as meaning "invoke this script using amm", and fire things up.

Hello, world

Every programming tutorial starts with "hello, world", so let's do our own slightly more interesting rendition. This was the first function I wrote in Ammonite, and I left it in there as a demonstration. In this case, it looks like:

@main
def hello() = {
  println(s"The API token is $layerApiToken")
  println(s"The App ID is $layerAppId")
}

Fairly ordinary Scala there, but notice that @main annotation at the top. That tells Ammonite that this is an entry point for the script.

Also, note the variables being interpolated in the printlns. Those are defined at the top of the script:

lazy val layerApiToken = sys.env("LAYER_API_TOKEN")
lazy val layerAppId = sys.env("LAYER_APP_ID")

These are shell environment variables that we are reading in and will be using in the script. (In this case, secret keys that shouldn't be checked into git, so we don't want them in the script proper. As a general rule, you should never include anything remotely secret in checked-in code: they should always be managed externally.)

If we only had one entry point, we could simply say layer.sc at the command line, and Ammonite would run that function. But since our file actually has a bunch of @main functions, we instead need to say which one to invoke:

layer.sc hello

The program will compile, run, and print its output:

The API token is [REDACTED]
The App ID is [REDACTED]

There's no overhead -- it's pretty much as easy as any other shell script.

Let's Get Started

The point of this script is to export stuff from Layer, and there are a bunch of steps to that. To make this manageable, we're going to write each step as a separate entry point, which can be invoked from the command line.

(We could write this as a bunch of tiny separate scripts instead, of course. But it's much easier to do it all in one file, so we can share code and keep everything together.)

So let's create an entry point for the first step:

@main
def start() = {
  ...
}

Like our hello() above, by adding @main, we declare that this is a top-level entry point, which we can invoke from the command line as:

layer.sc start

That needs to make a request to Layer, to ask it to generate a dump of our data, which looks like this:

  val response = requests.post(
    s"https://api.layer.com/apps/$layerAppId/exports",
    headers = standardHeaders
  ).text

What's going on here? Let's take it apart.

Ammonite comes bundled with some useful libraries. One of them, which is simply available to you for free, is requests-scala, a lightweight HTTP client ported from the Python Requests client.

Making an HTTP POST request is as easy as shown here: we just call requests.post() with the URL that we need to fetch from, and whatever parameters are needed. In this case, we aren't actually sending a request body (the exports route doesn't require one), but we do need some headers. Since those are common to all of the Layer calls, we've pulled them out into a common function:

def standardHeaders = Map(
  "Accept" -> "application/vnd.layer+json; version=3.0",
  "Authorization" -> s"Bearer $layerApiToken",
  "Content-Type" -> "application/json"
)

requests.post() is a straightforward blocking function -- not what you would want for a high-performance application, but totally appropriate for a shell script like this. By the time it returns, we have the response from Layer, and turn that into a String by calling .text on it.

Playing with JSON

The response from Layer is a JSON document, so our next step is to parse that. Fortunately, Ammonite also includes uPickle right out of the box, so parsing the JSON is a one-liner:

  val json = ujson.read(response).obj

(That's just reading in a raw JSON structure. If we wanted to deserialize it into a Scala data structure, we could also do that with uPickle, but for now we're going to keep it simple.)

The bit we need out of there is the status_url field, which tells us the URL that we will need to poll in order to find out when the export is finished:

  json.get("status_url").map(_.str) match {
    case Some(statusUrl) => {
      // This is the URL that we will ping in check():
      println(s"\nStatus URL: $statusUrl")
      write.over(statusUrlPath, statusUrl)
      // Bump the download counter:
      val downloadCounter: Int = getDownloadCounter()
      write.over(downloadCounterPath, (downloadCounter + 1).toString)
    }
    case None => {
      println(s"\nDidn't get a statusUrl; something went wrong here!")
    }
  }

This is all fairly straightforward Scala.json.get() fetches the given field, returning an Option[Value]; .map(_.str) interprets that Value as a String, which we put into statusUrl. As for write.over(), read on...

Working with Files

Since this process takes a long time, with several steps, we're going to store our interim values on the filesystem. This creates a file containing just the statusUrl:

      write.over(statusUrlPath, statusUrl)

We're now making use of Ammonite-Ops, which we pull in at the top of our file:

import ammonite.ops._

(Sadly, IntelliJ isn't smart about Ammonite and has a bad habit of deciding that this import isn't being used. If this happens, in IntelliJ go to:

Settings > Editor > Code Style > Scala > Imports always marked as used:

and add ammonite.ops._ there.)

This provides a notion of Path -- a way to get to a location on the filesystem. Up in the file, we have:

lazy val wd = pwd
lazy val statusUrlPath = wd / "currentStatusUrl.txt"

Nothing strange here: statusUrlPath is just a constant, giving the path to our file. And it shows an advantage of doing everything in one script: all of our steps are sharing the same vals, so it's easy to keep everything consistent between the steps.

Finally, since we want to keep a bunch of versions of the export (in case something goes wrong), we increment a counter saying which invocation of the process we are currently in:

      val downloadCounter: Int = getDownloadCounter()
      write.over(downloadCounterPath, (downloadCounter + 1).toString)

And getDownloadCounter() lets us see some basic file input:

def getDownloadCounter(): Int = {
  val downloadCounterStr = read.lines(downloadCounterPath).headOption.getOrElse(throw new Exception(s"Couldn't read downloadCounter!"))
  downloadCounterStr.toInt
}

read.lines() is another function from Ammonite-Ops -- as the name implies, it reads the specified file in as a Seq[String], one per line. From there, we take the first line, and parse it as an Int. All concise and easy.

Are We Done Yet?

Layer's export process, I found, could take anywhere from 2-10 hours. So we need to poll that periodically, which we will do with this check() command:

@main
def check() = {
    ...
}

Again, we simply invoke that as:

layer.sc check

That uses the status URL that we got earlier:

  val statusUrl = read! statusUrlPath

read! is another Ammonite-Ops function -- given a Path, it reads the contents of that file in as a String. We check that URL using requests-scala again:

  val response = requests.get(
    statusUrl,
    headers = standardHeaders
  ).text

This makes an HTTP GET request to that URL, again using the standard headers to authenticate, and uses .text to turn the response into a String.

We parse the JSON document and fetch the status field in the same way we did in start():

  val obj = ujson.read(response).obj
  val status = obj.getString("status")

getString() is an extension function that we added to make life a bit easier, showing that all of Scala (including fancy stuff like implicit classes) is available to us:

implicit class RichObj(obj: scala.collection.mutable.LinkedHashMap[String, ujson.Value]) {
  // Convenience function for getting a required string-valued field from a JSON object
  def getString(name: String): String = obj.get(name).map(_.str).getOrElse(throw new Exception(s"Couldn't find a $name!"))
}

Given that status, we do an ordinary match on it, and move on to the next step if we're ready:

  status match {
    case "pending" => println(s"Still waiting for this to complete...")
    case "completed" => {
      // Excellent -- grab the URL we will need for download():
      val downloadUrl = obj.getString("download_url")
      write.over(downloadUrlPath, downloadUrl)
      // Store the key and initialization vector that we will need in decode() for this archive:
      val downloadCounter = getDownloadCounter()
      val keyPath = aesKey(downloadCounter)
      val ivPath = aesIV(downloadCounter)
      val key = obj.getString("encrypted_aes_key")
      val iv = obj.getString("aes_iv")
      write.over(keyPath, key)
      write.over(ivPath, iv)
      println(s"We're ready to roll! Call the download function next.")
    }
    case other => println(s"Got unexpected status $other")
  }

The download_url is straightforward -- we're just grabbing the String and storing it in a file as seen above.

The aesKey and aesIV are slightly more involved: these are the decryption key and initialization vector that we will need in order to read the file once we've downloaded it. Since these are different for each download, we want to keep them in separate files, so the Path for each is based on the current download counter:

def aesKey(counter: Int) = wd / s"aesKey-$counter"
def aesIV(counter: Int) = wd / s"aesIV-$counter"

There's nothing complicated about constructing a Path -- we can use String interpolation the same way we always do in Scala. And once we have the Paths, we write them out using write.over() as before.

Downloading with curl

The download() entry point should be pretty straightforward by now:

@main
def download() = {
  val downloadCounter: Int = getDownloadCounter()
  val downloadUrl = read! downloadUrlPath
  val filename = downloadedFilePath(downloadCounter)
  implicit val cwd = wd
  %("curl", downloadUrl, "-o", filename)
  println("Done!")
}

That %() function is new, though. In this case, we want to download the file directly to the filesystem, not read it into a String, and the easiest way to do that is to use the operating system's curl command. That's what %() lets you do: call an outside program, and return once it is finished.

So this function figures out which download we're currently working on, fetches the URL that we saved in the previous step, and uses curl to fetch the actual archive from that URL. Nice and easy.

Encryption is Always the Hard Part

Of course, this archive is full of sensitive data, so it's not plaintext. Instead, it is encrypted using an encryption key and initialization vector (which we received two steps ago), and the encryption key is itself encrypted with our key.

When we want to decrypt the archive (potentially a much later point than the time when we downloaded it), we need a bunch of steps. In principle, we could just pipe between them; in practice, that wasn't as easy to do with Ammonite. (Showing how quickly the tool is continuing to improve, pipes have been added since I wrote this script.)

So instead, we're going to treat it as a bunch of individual steps.

To support all of this, we're going to add a convenient little utility function:

def runScript(script: String): Seq[String] = {
  implicit val cwd = wd
  %%("bash", "-c", script).out.lines
}

Again, this is an ordinary function, but we're introducing a new function, %%(). Where %(), which we used above, just fires up a shell and executes the command, %%() actually captures the output in a useful format. So here, we are taking a command line, executing that using bash, and returning the output lines.

The new decode entry point looks like this:

@main
def decode(version: Int = getDownloadCounter()) = {
    ...
}

Notice that, unlike our earlier entry points, this one takes a parameter. So we can call it like this:

layer.sc decode 4

to tell it to decode archive #4. (Common types like Int get parsed automatically for you.) Of course, we typically want to decode the last archive, so we default to the current value in the counter using ordinary Scala default parameter values.

The steps of decode() are a bit more complex, but at this point you should be able to read them. First, we read in the decryption key, and decrypt that against our own private key (which is in the local file layer-export-key.pem):

  val encryptedKey = read! aesKey(version)
  val key =
    runScript(
      s"""echo "$encryptedKey" | base64 --decode | openssl rsautl -decrypt -inkey layer-export-key.pem | hexdump -ve '1/1 "%.2x"'"""
    ).head

That produces one line of output, the key itself, so .head gives us what we are looking for.

Then we read in the initialization vector, which is just base64 encoded:

  val encodedIV = read! aesIV(version)
  val iv = runScript(s"""echo $encodedIV | base64 --decode | hexdump -ve '1/1 "%.2x"'""").head

Then we use those to actually decrypt the archive, which gives us a gzip'ed tarball:

  val targzPath = wd / s"download_$version.tar.gz"
  runScript(s"""openssl enc -in ${downloadedFilePath(version)} -out $targzPath -d -aes-256-cbc -K "$key" -iv "$iv" """)

Unzipping and untarring are relatively easy:

  runScript(s"""gunzip $targzPath""")
  val tarPath = wd / s"download_$version.tar"
  runScript(s"""tar -xvf $tarPath""")

Finally, we put the decrypted archive into an appropriate directory and clean up, using more of the convenient (and clear) tools provided by Ammonite:

  val outputDir = wd / s"decoded_$version"
  mkdir! outputDir
  mv(wd / "export.json", outputDir / "export.json")
  rm! tarPath

Conclusion: That Was Easier Than Expected

There's a lot of text above, but honestly, it took much less time to write this script than it did to write this article -- about four hours total, from a standing start of never having worked with Ammonite before. Knowing what I know now, knocking out future scripts should be a snap.

Are Scala Scripts going to take over the world? Probably not -- the strongly-typed view of things is very much a matter of taste, and most folks prefer quick-and-dirty scripting languages, which is usually fine.

But for those of you who already know Scala, and who consider shell scripting a chore, I strongly recommend checking out Ammonite. Frankly, I think it's a joy: it lets me do scripting using the nice reliable language that I already know and love, and the result is reliable scripts that I know are going to work the way I want.

We hope you found this article helpful but remember that you are solely responsible for your development. More details and disclaimers here.

Mark "Justin" Waks