
from transforms.base import *
import struct


class Lsb(Transform):
    """
    Keep only the <num_bits> lsb of each byte
    """
    category = "stegano"
    name = "lsb"
    icon = "wxART_FILETYPE_IMAGE"

    def run(self, data:bytes, num_bits:int=1, start_bit:int=0):
        if start_bit >= num_bits:
            raise ValueError("start_bit must be < num_bits")
        res = bytearray()
        curval = 0
        curi = 0
        first = True
        for c in data:
            if first:
                first = False
                s = start_bit
            else:
                 s = 0
            for i in range(s, num_bits):
                if c & (1 << i):
                    curval = curval | (1 << curi)
                curi += 1
                if curi >= 8:
                    res.append(curval)
                    curval = 0
                    curi = 0
        return res
