#!/usr/bin/env python

from __future__ import print_function

"""Simple PEM encode, decode and extraction
"""

"""
RCSid:
	$Id: pem.py,v 1.9 2023/04/25 16:54:17 sjg Exp $

	@(#) Copyright (c) 2012-2021 Simon J. Gerraty

	This file is provided in the hope that it will
	be of use.  There is absolutely NO WARRANTY.
	Permission to copy, redistribute or otherwise
	use this file is hereby granted provided that 
	the above copyright notice and this notice are
	left intact. 
      
	Please send copies of changes and bug-fixes to:
	sjg@crufty.net

"""

import base64
from p3compat import *

def pem_encode(id, data):
    """PEM encode data with given id"""
    pem = ['-----BEGIN {0}-----'.format(id)]
    b64 = base64.b64encode(b(data))
    x = 0
    n = len(b64)
    while x < n:
        y = x + min(n - x, 64)
        pem.append(s(b64[x:y]))
        x = y
        
    pem.append('-----END {0}-----\n'.format(id))
    return '\n'.join(pem)

def pem_extract(id, lines, debug=0):
    """return the portion of lines which contains id"""
    s = -1
    for i in range(0, len(lines)):
        if s >= 0:
            if lines[i].startswith('-----END'):
                if debug:
                    print(i,lines[i])
                return lines[s:i+1]
            
        if not lines[i].startswith('-----BEGIN'):
            continue
        if lines[i].find(id) < 0:
            continue
        s = i
        if debug:
            print(i,lines[i])

    return None

def pem_decode(text):
    """extract and decode pem data"""
    # pem = data between the BEGIN and END markers
    text = s(text)
    begin = text.find('-----BEGIN')
    if begin >= 0:
        begin = text.find('\n', begin) + 1
        if begin > 0:
            pem = text[begin:]
            end = pem.find('-----END')
            if end > 0:
                pem = pem[:end]
                x = pem.find('\n\n')
                if x > 0:
                    pem = pem[x+2:]
                pem = pem.replace('\n', '')
            else:
                return None
        else:
            return None
    else:
        return None

    data = base64.b64decode(pem)
    return s(data)

if __name__ == '__main__':
    m0 = """
    Hello, World!
    Life! don't talk to me about life
    Here I am, brain the size of a planet...
    and what do they have me do?"""

    print(m0)
    pem = pem_encode('MESSAGE', m0)
    print(pem)
    m1 = pem_decode(pem)
    print(m1)
    m2 = pem_decode('junk\n' + pem + 'more junk\n')
    print(m2)
    
    assert m2 == m0

