<commit_before>import numpy as np

from errors import *
from diskimages import DiskImageBase
from segments import EmptySegment, ObjSegment, RawSectorsSegment, DefaultSegment, SegmentSaver
from utils import to_numpy

import logging
log = logging.getLogger(__name__)


class AtariDosDirent(object):
    # ATR Dirent structure described at http://atari.kensclassics.org/dos.htm
    format = np.dtype([
        ('FLAG', 'u1'),
        ('COUNT', '<u2'),
        ('START', '<u2'),
        ('NAME','S8'),
        ('EXT','S3'),
        ])

    def __init__(self, image, file_num=0, bytes=None):
        self.file_num = file_num
        self.flag = 0
        self.opened_output = False
        self.dos_2 = False
        self.mydos = False
        self.is_dir = False
        self.locked = False
        self.in_use = False
        self.deleted = False
        self.num_sectors = 0
        self.starting_sector = 0
        self.filename = ""
        self.ext = ""
        self.is_sane = True
        self.current_sector = 0
        self.current_read = 0
        self.sectors_seen = None
        self.parse_raw_dirent(image, bytes)
    
    def __str__(self):
        flags = self.summary()
        return "File #%-2d (%s) %03d %-8s%-3s  %03d" % (self.file_num, flags, self.starting_sector, self.filename, self.ext, self.num_sectors)
    
    def summary(self):
        output = "o" if self.opened_output else "."
        dos2 = "2" if self.dos_2 else "."
        mydos = "m" if self.mydos else "."
        in_use = "u" if self.in_use else "."
        deleted = "d" if self.deleted else "."
        locked = "*" if self.locked else " "
        flags = "%s%s%s%s%s%s" % (output, dos2, mydos, in_use, deleted, locked)
        return flags
    
    @property
    def verbose_info(self):
        flags = []
        if self.opened_output: flags.append("OUT")
        if self.dos_2: flags.append("DOS2")
        if self.mydos: flags.append("MYDOS")
        if self.in_use: flags.append("IN_USE")
        if self.deleted: flags.append("DEL")
        if self.locked: flags.append("LOCK")
        return "flags=[%s]" % ", ".join(flags)
    
    def parse_raw_dirent(self, image, bytes):
        if bytes is None:
            return
        values = bytes.view(dtype=self.format)[0]  
        flag = values[0]
        self.flag = flag
        self.opened_output = (flag&0x01) > 0
        self.dos_2 = (flag&0x02) > 0
        self.mydos = (flag&0x04) > 0
        self.is_dir = (flag&0x10) > 0
        self.locked = (flag&0x20) > 0
        self.in_use = (flag&0x40) > 0
        self.deleted = (flag&0x80) > 0
        self.num_sectors = int(values[1])
        self.starting_sector = int(values[2])
        self.filename = str(values[3]).rstrip()
        self.ext = str(values[4]).rstrip()
        self.is_sane = self.sanity_check(image)
    
    def sanity_check(self, image):
        if not self.in_use:
            return True
        if not image.header.sector_is_valid(self.starting_sector):
            return False
        if self.num_sectors < 0 or self.num_sectors > image.header.max_sectors:
            return False
        return True
    
    def start_read(self, image):
        if not self.is_sane:
            raise InvalidDirent("Invalid directory entry '%s'" % str(self))
        self.current_sector = self.starting_sector
        self.current_read = self.num_sectors
        self.sectors_seen = set()
    
    def read_sector(self, image):
        raw, pos, size = image.get_raw_bytes(self.current_sector)
        bytes, num_data_bytes = self.process_raw_sector(image, raw)
        return bytes, self.current_sector == 0, pos, num_data_bytes

    def process_raw_sector(self, image, raw):
        file_num = raw[-3] >> 2
        if file_num != self.file_num:
            raise FileNumberMismatchError164()
        self.sectors_seen.add(self.current_sector)
        next_sector = ((raw[-3] & 0x3) << 8) + raw[-2]
        if next_sector in self.sectors_seen:
            raise InvalidFile("Bad sector pointer data: attempting to reread sector %d" % next_sector)
        self.current_sector = next_sector
        num_bytes = raw[-1]
        return raw[0:num_bytes], num_bytes
    
    def get_filename(self):
        ext = ("." + self.ext) if self.ext else ""
        return self.filename + ext


class MydosDirent(AtariDosDirent):
    def process_raw_sector(self, image, raw):
        # No file number stored in the sector data; two full bytes available
        # for next sector
        self.current_sector = (raw[-3] << 8) + raw[-2]
        num_bytes = raw[-1]
        return raw[0:num_bytes], num_bytes


class XexSegmentSaver(SegmentSaver):
    export_data_name = "Atari 8-bit Executable"
    export_extensions = [".xex"]


class XexSegment(ObjSegment):
    savers = [SegmentSaver, XexSegmentSaver]


class AtariDosFile(object):
    """Parse a binary chunk into segments according to the Atari DOS object
    file format.
    
    Ref: http://www.atarimax.com/jindroush.atari.org/afmtexe.html
    """
    def __init__(self, rawdata):
        self.rawdata = rawdata
        self.size = len(rawdata)
        self.segments = []
        self.parse_segments()
    
    def __str__(self):
        return "\n".join(str(s) for s in self.segments) + "\n"
    
    def parse_segments(self):
        r = self.rawdata
        b = r.get_data()
        pos = 0
        first = True
        log.debug("Initial parsing: size=%d" % self.size)
        while pos < self.size:
            if pos + 1 < self.size:
                header, = b[pos:pos+2].view(dtype='<u2')
            else:
                self.segments.append(ObjSegment(r[pos:pos + 1], pos, pos + 1, 0, 1, "Incomplete Data"))
                break
            log.debug("header parsing: header=0x%x" % header)
            if header == 0xffff:
                # Apparently 0xffff header can appear in any segment, not just
                # the first.  Regardless, it is ignored everywhere.
                pos += 2
            elif first:
                raise InvalidBinaryFile
            first = False
            if len(b[pos:pos + 4]) < 4:
                self.segments.append(ObjSegment(r[pos:pos + 4], 0, 0, "Short Segment Header"))<commit_msg>Fix for incorrect ObjSegment call when there's a short segment header<commit_after>import numpy as np

from errors import *
from diskimages import DiskImageBase
from segments import EmptySegment, ObjSegment, RawSectorsSegment, DefaultSegment, SegmentSaver
from utils import to_numpy

import logging
log = logging.getLogger(__name__)


class AtariDosDirent(object):
    # ATR Dirent structure described at http://atari.kensclassics.org/dos.htm
    format = np.dtype([
        ('FLAG', 'u1'),
        ('COUNT', '<u2'),
        ('START', '<u2'),
        ('NAME','S8'),
        ('EXT','S3'),
        ])

    def __init__(self, image, file_num=0, bytes=None):
        self.file_num = file_num
        self.flag = 0
        self.opened_output = False
        self.dos_2 = False
        self.mydos = False
        self.is_dir = False
        self.locked = False
        self.in_use = False
        self.deleted = False
        self.num_sectors = 0
        self.starting_sector = 0
        self.filename = ""
        self.ext = ""
        self.is_sane = True
        self.current_sector = 0
        self.current_read = 0
        self.sectors_seen = None
        self.parse_raw_dirent(image, bytes)
    
    def __str__(self):
        flags = self.summary()
        return "File #%-2d (%s) %03d %-8s%-3s  %03d" % (self.file_num, flags, self.starting_sector, self.filename, self.ext, self.num_sectors)
    
    def summary(self):
        output = "o" if self.opened_output else "."
        dos2 = "2" if self.dos_2 else "."
        mydos = "m" if self.mydos else "."
        in_use = "u" if self.in_use else "."
        deleted = "d" if self.deleted else "."
        locked = "*" if self.locked else " "
        flags = "%s%s%s%s%s%s" % (output, dos2, mydos, in_use, deleted, locked)
        return flags
    
    @property
    def verbose_info(self):
        flags = []
        if self.opened_output: flags.append("OUT")
        if self.dos_2: flags.append("DOS2")
        if self.mydos: flags.append("MYDOS")
        if self.in_use: flags.append("IN_USE")
        if self.deleted: flags.append("DEL")
        if self.locked: flags.append("LOCK")
        return "flags=[%s]" % ", ".join(flags)
    
    def parse_raw_dirent(self, image, bytes):
        if bytes is None:
            return
        values = bytes.view(dtype=self.format)[0]  
        flag = values[0]
        self.flag = flag
        self.opened_output = (flag&0x01) > 0
        self.dos_2 = (flag&0x02) > 0
        self.mydos = (flag&0x04) > 0
        self.is_dir = (flag&0x10) > 0
        self.locked = (flag&0x20) > 0
        self.in_use = (flag&0x40) > 0
        self.deleted = (flag&0x80) > 0
        self.num_sectors = int(values[1])
        self.starting_sector = int(values[2])
        self.filename = str(values[3]).rstrip()
        self.ext = str(values[4]).rstrip()
        self.is_sane = self.sanity_check(image)
    
    def sanity_check(self, image):
        if not self.in_use:
            return True
        if not image.header.sector_is_valid(self.starting_sector):
            return False
        if self.num_sectors < 0 or self.num_sectors > image.header.max_sectors:
            return False
        return True
    
    def start_read(self, image):
        if not self.is_sane:
            raise InvalidDirent("Invalid directory entry '%s'" % str(self))
        self.current_sector = self.starting_sector
        self.current_read = self.num_sectors
        self.sectors_seen = set()
    
    def read_sector(self, image):
        raw, pos, size = image.get_raw_bytes(self.current_sector)
        bytes, num_data_bytes = self.process_raw_sector(image, raw)
        return bytes, self.current_sector == 0, pos, num_data_bytes

    def process_raw_sector(self, image, raw):
        file_num = raw[-3] >> 2
        if file_num != self.file_num:
            raise FileNumberMismatchError164()
        self.sectors_seen.add(self.current_sector)
        next_sector = ((raw[-3] & 0x3) << 8) + raw[-2]
        if next_sector in self.sectors_seen:
            raise InvalidFile("Bad sector pointer data: attempting to reread sector %d" % next_sector)
        self.current_sector = next_sector
        num_bytes = raw[-1]
        return raw[0:num_bytes], num_bytes
    
    def get_filename(self):
        ext = ("." + self.ext) if self.ext else ""
        return self.filename + ext


class MydosDirent(AtariDosDirent):
    def process_raw_sector(self, image, raw):
        # No file number stored in the sector data; two full bytes available
        # for next sector
        self.current_sector = (raw[-3] << 8) + raw[-2]
        num_bytes = raw[-1]
        return raw[0:num_bytes], num_bytes


class XexSegmentSaver(SegmentSaver):
    export_data_name = "Atari 8-bit Executable"
    export_extensions = [".xex"]


class XexSegment(ObjSegment):
    savers = [SegmentSaver, XexSegmentSaver]


class AtariDosFile(object):
    """Parse a binary chunk into segments according to the Atari DOS object
    file format.
    
    Ref: http://www.atarimax.com/jindroush.atari.org/afmtexe.html
    """
    def __init__(self, rawdata):
        self.rawdata = rawdata
        self.size = len(rawdata)
        self.segments = []
        self.parse_segments()
    
    def __str__(self):
        return "\n".join(str(s) for s in self.segments) + "\n"
    
    def parse_segments(self):
        r = self.rawdata
        b = r.get_data()
        pos = 0
        first = True
        log.debug("Initial parsing: size=%d" % self.size)
        while pos < self.size:
            if pos + 1 < self.size:
                header, = b[pos:pos+2].view(dtype='<u2')
            else:
                self.segments.append(ObjSegment(r[pos:pos + 1], pos, pos + 1, 0, 1, "Incomplete Data"))
                break
            log.debug("header parsing: header=0x%x" % header)
            if header == 0xffff:
                # Apparently 0xffff header can appear in any segment, not just
                # the first.  Regardless, it is ignored everywhere.
                pos += 2
            elif first:
                raise InvalidBinaryFile
            first = False
            if len(b[pos:pos + 4]) < 4:
                self.segments.append(ObjSegment(r[pos:pos + 4], 0, 0, 0, len(b[pos:pos + 4]), "Short Segment Header"))
{'fraction_non_alphanumeric': 0.11717998075072185, 'fraction_numerical': 0.015720243824189926, 'mean_word_length': 2.144766708701135, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6157529', 'n_tokens_mistral': 4193, 'n_tokens_neox': 3959, 'n_words': 966}