capafmt.py

# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Reformat the given capa rule into a consistent style.
Use the -i flag to update the rule in-place.

Usage:

   $ python capafmt.py -i foo.yml
"""

import sys
import logging
import argparse
from pathlib import Path

import capa.main
import capa.rules

logger = logging.getLogger("capafmt")


def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]

    parser = argparse.ArgumentParser(description="Capa rule formatter.")
    capa.main.install_common_args(parser)
    parser.add_argument("path", type=str, help="Path to rule to format")
    parser.add_argument(
        "-i",
        "--in-place",
        action="store_true",
        dest="in_place",
        help="Format the rule in place, otherwise, write formatted rule to STDOUT",
    )
    parser.add_argument(
        "-c",
        "--check",
        action="store_true",
        help="Don't output (reformatted) rule, only return status. 0 = no changes, 1 = would reformat",
    )
    args = parser.parse_args(args=argv)

    try:
        capa.main.handle_common_args(args)
    except capa.main.ShouldExitError as e:
        return e.status_code

    rule = capa.rules.Rule.from_yaml_file(args.path, use_ruamel=True)
    reformatted_rule = rule.to_yaml()

    if args.check:
        if rule.definition == reformatted_rule:
            logger.info("rule is formatted correctly, nice! (%s)", rule.name)
            return 0
        else:
            logger.info("rule requires reformatting (%s)", rule.name)
            if "\r\n" in rule.definition:
                logger.info("please make sure that the file uses LF (\\n) line endings only")
            return 1

    if args.in_place:
        Path(args.path).write_bytes(reformatted_rule.encode("utf-8"))
    else:
        print(reformatted_rule)

    return 0


if __name__ == "__main__":
    sys.exit(main())