check_key.py

#!/usr/bin/env python3

# http://sonelli.freshdesk.com/support/solutions/articles/139632-juicessh-supported-private-key-formats-openssh-pem-
# don't use this program
# use: ssh-keygen -t rsa -b 4096 -f you_key_file -m pem

"""
Retval   0: File is ok
Retval -10: File is too big
Retval -20: Header not found
"""
import sys
import logging
from pathlib import Path
from argparse import ArgumentParser


def check_key(file, log):
    """
    Returns  0 if the file is ok.
    Returns 10 if the file is too big.
    Returns 20 if the header was not found.
    """
    headers = [
        "-----{} RSA PRIVATE KEY-----",
        "-----{} DSA PRIVATE KEY-----",
        "-----{} PRIVATE KEY-----",
        "-----{} ENCRYPTED PRIVATE KEY-----",
    ]
    headers.sort(key=len)
    fsize = file.stat().st_size
    if fsize < 8 * 1024 ** 1:
        with file.open() as fd:
            header, *content, footer = fd.read().strip().splitlines()
        header_ok = any(header.startswith(h.format('BEGIN')) for h in headers)
        footer_ok = any(footer.startswith(h.format('END')) for h in headers)
        hf_err = False
        if header_ok and footer_ok:
            log.info('File is ok.')
            return 0
        if not header_ok:
            log.info('Header is not ok')
            log.debug(header)
            hf_err = True
        if not footer_ok:
            log.info('Footer is not ok')
            log.debug(footer)
            hf_err = True
        if hf_err:
            return 20
    log.info('File is too big')
    log.debug(f'File {file.name} has a size of {fsize // 1024**1:.2f} kiB.')
    return 10


def fix(file):
    """
    Dumb function to replace header and footer of the file
    and replaces it with the correct header
    """
    content = file.read_bytes().strip().splitlines()[1:-1]
    header = b"-----BEGIN RSA PRIVATE KEY-----"
    footer = b"-----END RSA PRIVATE KEY-----"
    new_content = [header, *content, footer]
    file.write_bytes(b'\n'.join(new_content))


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument('file', type=Path, help='SSH Private Keyfile')
    parser.add_argument('-i', action='store_true', help='Info')
    parser.add_argument('-v', action='store_true', help='Verbose')
    parser.add_argument('-f', action='store_true', help='Fix file inplace, expecting RSA')
    args = parser.parse_args()
    logging.basicConfig()
    log = logging.getLogger(sys.argv[0])
    if args.i:
        log.setLevel(logging.INFO)
    if args.v:
        log.setLevel(logging.DEBUG)
    retval = check_key(args.file, log)
    if args.f and retval != 0:
        log.info('Fixing file')
        fix(args.file)
    sys.exit(retval)


if __name__ == '__main__':
    main()