<reponame>simonowen/mgtdisklib
# Image wrapper for disk image containers (MGT/SAD/EDSK).
#
# Part of https://github.com/simonowen/mgtdisklib

import os
import gzip
from typing import Optional


class Image:
    def __init__(self, *, spt: int = 10) -> None:
        self.path: Optional[str] = None
        self.spt: int = spt
        self.compressed: bool = False
        self.data: bytearray = bytearray(80 * 2 * self.spt * 512)

    @staticmethod
    def open(path: str):
        """Create Image object from disk image file"""
        with open(path, 'rb') as f:
            data = bytearray(f.read())
            if data[:2] == b'\x1f\x8b':
                with gzip.open(path, 'rb') as f2:
                    data = bytearray(f2.read())
                    compressed = True
            else:
                compressed = False

        image: Optional[Image] = None
        if len(data) == 819200 or len(data) == 737280:
            image = MGTImage()
            image.spt = len(data) // (80 * 2 * 512)
        elif len(data) == 819222 or len(data) == 737302:
            image = SADImage()
            image.spt = (len(data) - 22) // (80 * 2 * 512)
        elif len(data) == 860416 or len(data) == 778496:
            image = EDSKImage()
            image.spt = (((len(data) - 0x100) // (80 * 2)) - 0x100) // 512
        else:
            raise RuntimeError(f'{path} is not a supported disk image')

        image.path = os.path.abspath(path)
        image.data = data
        image.compressed = compressed
        return image

    def save(self, path: str, *, compressed: bool = False) -> None:
        """Save disk data to image file"""
        if compressed:
            with gzip.open(path, 'wb') as f:
                f.write(self.data)
        else:
            with open(path, 'wb') as f2:
                f2.write(self.data)

    def read_sector(self, track: int, sector: int) -> bytes:
        """Read sector data for given sector location"""

        offset = self.sector_offset(track, sector)
        return bytes(self.data[offset:offset+512])

    def write_sector(self, track: int, sector: int, data: bytes) -> None:
        """Write sector data to given sector location"""

        if len(data) != 512:
            raise ValueError('sector data should be 512 bytes')

        offset = self.sector_offset(track, sector)
        self.data[offset:offset+512] = data

    def sector_offset(self, track: int, sector: int) -> int:
        """Calculate sector data offset in image data"""

        if track < 0 or (track & 0x7f) >= 80 or sector < 1 or sector > self.spt:
            raise ValueError(f'invalid sector location: track {track} sector {sector}')

        return ((track & 0x7f) * self.spt * 2 + (track >> 7) * self.spt + (sector - 1)) * 512


class MGTImage(Image):
    pass


class SADImage(Image):
    def sector_offset(self, track: int, sector: int) -> int:
        """Calculate sector data offset in SAD image"""

        if track < 0 or (track & 0x7f) >= 80 or sector < 1 or sector > self.spt:
            raise ValueError(f'invalid sector location track {track} sector {sector}')

        return 22 + (((track >> 7) * 80 + (track & 0x7f)) * self.spt + (sector - 1)) * 512


class EDSKImage(Image):
    def sector_offset(self, track: int, sector: int) -> int:
        """Calculate sector data offset in EDSK image"""

        if track < 0 or (track & 0x7f) >= 80 or sector < 1 or sector > self.spt:
            raise ValueError(f'invalid sector location track {track} sector {sector}')

        track_offset = 0x100 + ((track & 0x7f) * 2 + (track >> 7)) * (0x100 + self.spt * 512)
        sectors = [self.data[track_offset + 0x1a + i * 8] for i in range(self.spt)]
        return track_offset + 0x100 + sectors.index(sector) * 512

# Image wrapper for disk image containers (MGT/SAD/EDSK).
#
# Part of https://github.com/simonowen/mgtdisklib

import os
import gzip
from typing import Optional


class Image:
    def __init__(self, *, spt: int = 10) -> None:
        self.path: Optional[str] = None
        self.spt: int = spt
        self.compressed: bool = False
        self.data: bytearray = bytearray(80 * 2 * self.spt * 512)

    @staticmethod
    def open(path: str):
        """Create Image object from disk image file"""
        with open(path, 'rb') as f:
            data = bytearray(f.read())
            if data[:2] == b'\x1f\x8b':
                with gzip.open(path, 'rb') as f2:
                    data = bytearray(f2.read())
                    compressed = True
            else:
                compressed = False

        image: Optional[Image] = None
        if len(data) == 819200 or len(data) == 737280:
            image = MGTImage()
            image.spt = len(data) // (80 * 2 * 512)
        elif len(data) == 819222 or len(data) == 737302:
            image = SADImage()
            image.spt = (len(data) - 22) // (80 * 2 * 512)
        elif len(data) == 860416 or len(data) == 778496:
            image = EDSKImage()
            image.spt = (((len(data) - 0x100) // (80 * 2)) - 0x100) // 512
        else:
            raise RuntimeError(f'{path} is not a supported disk image')

        image.path = os.path.abspath(path)
        image.data = data
        image.compressed = compressed
        return image

    def save(self, path: str, *, compressed: bool = False) -> None:
        """Save disk data to image file"""
        if compressed:
            with gzip.open(path, 'wb') as f:
                f.write(self.data)
        else:
            with open(path, 'wb') as f2:
                f2.write(self.data)

    def read_sector(self, track: int, sector: int) -> bytes:
        """Read sector data for given sector location"""

        offset = self.sector_offset(track, sector)
        return bytes(self.data[offset:offset+512])

    def write_sector(self, track: int, sector: int, data: bytes) -> None:
        """Write sector data to given sector location"""

        if len(data) != 512:
            raise ValueError('sector data should be 512 bytes')

        offset = self.sector_offset(track, sector)
        self.data[offset:offset+512] = data

    def sector_offset(self, track: int, sector: int) -> int:
        """Calculate sector data offset in image data"""

        if track < 0 or (track & 0x7f) >= 80 or sector < 1 or sector > self.spt:
            raise ValueError(f'invalid sector location: track {track} sector {sector}')

        return ((track & 0x7f) * self.spt * 2 + (track >> 7) * self.spt + (sector - 1)) * 512


class MGTImage(Image):
    pass


class SADImage(Image):
    def sector_offset(self, track: int, sector: int) -> int:
        """Calculate sector data offset in SAD image"""

        if track < 0 or (track & 0x7f) >= 80 or sector < 1 or sector > self.spt:
            raise ValueError(f'invalid sector location track {track} sector {sector}')

        return 22 + (((track >> 7) * 80 + (track & 0x7f)) * self.spt + (sector - 1)) * 512


class EDSKImage(Image):
    def sector_offset(self, track: int, sector: int) -> int:
        """Calculate sector data offset in EDSK image"""

        if track < 0 or (track & 0x7f) >= 80 or sector < 1 or sector > self.spt:
            raise ValueError(f'invalid sector location track {track} sector {sector}')

        track_offset = 0x100 + ((track & 0x7f) * 2 + (track >> 7)) * (0x100 + self.spt * 512)
        sectors = [self.data[track_offset + 0x1a + i * 8] for i in range(self.spt)]
        return track_offset + 0x100 + sectors.index(sector) * 512
# Image wrapper for disk image containers (MGT/SAD/EDSK). # # Part of https://github.com/simonowen/mgtdisklib Create Image object from disk image file Save disk data to image file Read sector data for given sector location Write sector data to given sector location Calculate sector data offset in image data Calculate sector data offset in SAD image Calculate sector data offset in EDSK image