2

I'm trying to create a custom logger class for printing the log and save it to a file as encrypted at the same time. I used this reference. Here is my code:

import base64
import logging
from pprint import pprint
from Cryptodome.Cipher import AES
from Cryptodome.Hash import SHA256
from Cryptodome.Hash import MD5
from Cryptodome import Random


class logger:
    """
    Encrypt log messages to file as encrypted
    """
    class EncryptedLogFormatter(logging.Formatter):
        def __init__(self, key, fmt=None, datefmt=None):
            self._key = self.hash_gen(key, 16)
            super(logger.EncryptedLogFormatter, self).__init__(fmt=fmt, datefmt=datefmt)

        @staticmethod
        def hash_gen(key, size):
            """
            return a hash object of key base on size
            """
            key = MD5.new(key.encode('utf-8')).digest()  # use SHA-256 for a proper-sized AES key
            return key[:size]

        def format(self, record):
            # pprint(vars(record))
            message = record.msg  # log message to encrypt, if any
            asctime = record.asctime  # asctime to encrypt
            levelname = record.levelname  # levelname to encrypt

            if message:  # no sense to encrypt empty log messages
                iv = Random.new().read(AES.block_size)  # we'll be using CBC so generate an IV
                cipher = AES.new(self._key, AES.MODE_CBC, iv)
                # AES demands all blocks to be of `AES.block_size` so we have to pad the message
                # you can use any padding you prefer, I think PKCS#7 is the best option
                padding = AES.block_size - len(message) % AES.block_size
                # pad the message...
                message += chr(padding) * padding
                message_enc = iv + cipher.encrypt(message.encode())  # add iv and encrypt
                # finally, replace our plain-text message with base64 encoded encrypted one
                record.msg = base64.b64encode(message_enc).decode()

            if asctime:
                iv = Random.new().read(AES.block_size)
                cipher = AES.new(self._key, AES.MODE_CBC, iv)
                padding = AES.block_size - len(asctime) % AES.block_size
                asctime += chr(padding) * padding
                asctime_enc = iv + cipher.encrypt(asctime.encode())
                record.asctime = base64.b64encode(asctime_enc).decode()

            if levelname:
                iv = Random.new().read(AES.block_size)
                cipher = AES.new(self._key, AES.MODE_CBC, iv)
                padding = AES.block_size - len(levelname) % AES.block_size
                levelname += chr(padding) * padding
                levelname_enc = iv + cipher.encrypt(levelname.encode())
                record.levelname = base64.b64encode(levelname_enc).decode()

            return super(logger.EncryptedLogFormatter, self).format(record)

    def __init__(self, key, filename, level=logging.INFO, fmt='%(asctime)s:%(levelname)s: %(message)s', datefmt="%Y-%m-%d %H:%M:%S"):
        root = logging.getLogger()
        root.setLevel(level)
        ch = logging.StreamHandler()
        fh = logging.FileHandler(filename)
        formatter = logging.Formatter(fmt=fmt, datefmt=datefmt)
        ch.setFormatter(formatter)
        fh.setFormatter(logger.EncryptedLogFormatter(key, fmt, datefmt))
        root.addHandler(ch)
        root.addHandler(fh)

    def print(self, message):
        logging.info(message)


if __name__ == "__main__":
    logg = logger("abcdefg", 'Some path')
    logg.print("Hello")

Console output:

2018-08-12 13:21:07:INFO: Hello

File output:

2018-08-12 13:21:07:QcMrG7d7gvxwiagidFozC2v4kQukgnbXv5Hs2rMDAZQ=: Px4ZlIE7usOTTtbURDjrGW4VBXaIKH/F3vhs9pj5G3o=

It seems that the asctime hasn't been encrypted. What I want is to just use the user format and encrypt time, level and message. It would be better to just create the whole line encrypted but I don't know how to create the custom message for user input format.

4
  • What debugging have you done? One question would be: is it that asctime isn't set so the encryption is never done, or is it that the encrypted version is replaced by the parent? Have you considered getting the record from the parent then encrypting it? Commented Aug 12, 2018 at 9:16
  • 1
    you need to either encrypt the finale string after calling super().format() (and not before, since before that point the date isn't in string form and thus can't be encrypted) or just dont call super().format() at all and format the message youself Commented Aug 12, 2018 at 11:48
  • 1
    to clarify, you can't actually change asctime since it is determined by record.created every time you call format() Commented Aug 12, 2018 at 12:00
  • @jonrsharpe Thanks, with help of @AntiMatterDynamite I just encrypted the message before super().format() it was easier to implement. Commented Aug 13, 2018 at 4:43

1 Answer 1

2

As @AntiMatterDynamite said, the whole message can be encrypted before super().format():

import base64
import logging
from pprint import pprint
from Cryptodome.Cipher import AES
from Cryptodome.Hash import SHA256
from Cryptodome.Hash import MD5
from Cryptodome import Random


class logger:
    """
    Encrypt log messages to file as encrypted
    """
    class EncryptedLogFormatter(logging.Formatter):
        def __init__(self, key, fmt=None, datefmt=None):
            self._key = self.hash_gen(key, 16)
            super(logger.EncryptedLogFormatter, self).__init__(fmt=fmt, datefmt=datefmt)

        @staticmethod
        def hash_gen(key, size):
            """
            return a hash object of key base on size
            """
            key = MD5.new(key.encode('utf-8')).digest()  # use SHA-256 for a proper-sized AES key
            return key[:size]

        def format(self, record):
            # encrypt whole message instead of record.msg
            message = super().format(record)
            if message:  # no sense to encrypt empty log messages
                iv = Random.new().read(AES.block_size)  # we'll be using CBC so generate an IV
                cipher = AES.new(self._key, AES.MODE_CBC, iv)
                # AES demands all blocks to be of `AES.block_size` so we have to pad the message
                # you can use any padding you prefer, I think PKCS#7 is the best option
                padding = AES.block_size - len(message) % AES.block_size
                # pad the message...
                message += chr(padding) * padding
                message_enc = iv + cipher.encrypt(message.encode())  # add iv and encrypt
                # finally, replace our plain-text message with base64 encoded encrypted one
                return base64.b64encode(message_enc).decode()

    def __init__(self, key, filename, level=logging.INFO, fmt='%(asctime)s:%(levelname)s: %(message)s', datefmt="%Y-%m-%d %H:%M:%S"):
        root = logging.getLogger()
        root.setLevel(level)
        ch = logging.StreamHandler()
        fh = logging.FileHandler(filename)
        formatter = logging.Formatter(fmt=fmt, datefmt=datefmt)
        ch.setFormatter(formatter)
        fh.setFormatter(logger.EncryptedLogFormatter(key, fmt, datefmt))
        root.addHandler(ch)
        root.addHandler(fh)

    def print(self, message):
        logging.info(message)


if __name__ == "__main__":
    log = logger("abcdefg", 'Some path')
    log.print("Hello")
Sign up to request clarification or add additional context in comments.

Comments

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.