Sanitize a Pebble.state file, removing notices that refer to changes that do not exist

pebble_sanitize.py

import argparse
import json
import sys


def read_state(path):
    if path == '-':
        state = json.load(sys.stdin)
    else:
        with open(path) as f:
            state = json.load(f)
    return state

def check_consistency(state):
    # Changes is an object, where the attribute is the change id.
    changes = state['changes']
    for k, c in changes.items():
        if k != c.get('id'):
            print("change {} doesn't match {}\n".format(k, c), file=sys.stderr)
    # Notices is a list, with an id attribute, but the "key" of a
    # "change-update" notice is the change id
    notices = state['notices']
    badreference = []
    valid = []
    for n in notices:
        id_ = n.get('id')
        if id_ is None:
            # Invalid because it doesn't have an ID
            continue
        key = n.get("key")
        if key is None:
            # all notices should have a key
            continue
        typ = n.get("type")
        if typ == "change-update" and key not in changes:
            badreference.append(id_)
        else:
            valid.append(n)
    print("found {} missing references".format(len(badreference)), file=sys.stderr)
    return valid


def main(args):
    p = argparse.ArgumentParser("pebble_sanitize")
    p.add_argument("pebble_state", help="location of the pebble state file", default="/charm/container/pebble/.pebble.state")
    p.add_argument("-o", "--output", help="output file instead of stdout", default=None)
    opts = p.parse_args(args)
    state = read_state(opts.pebble_state)
    valid_notices = check_consistency(state)
    state['notices'] = valid_notices
    output = sys.stdout
    close = False
    try:
        if opts.output is not None and opts.output != "-":
            output = open(opts.output, "w")
            close = True
        json.dump(state, output)
    finally:
        if close:
            output.close()

if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))