update-extents.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303
  1. #!/usr/bin/env python3
  2. import array
  3. import struct
  4. def read_struct(fmt,buf,offset):
  5. out, = struct.unpack_from(fmt,buf,offset)
  6. return out, offset + struct.calcsize(fmt)
  7. class ISO(object):
  8. def __init__(self, path):
  9. with open(path, 'rb') as f:
  10. tmp = f.read()
  11. self.data = array.array('b', tmp)
  12. self.sector_size = 2048
  13. o = 0x10 * self.sector_size
  14. self.type, o = read_struct('B',self.data,o)
  15. self.id, o = read_struct('5s',self.data,o)
  16. self.version, o = read_struct('B',self.data,o)
  17. _unused0, o = read_struct('B',self.data,o)
  18. self.system_id, o = read_struct('32s',self.data,o)
  19. self.volume_id, o = read_struct('32s',self.data,o)
  20. _unused1, o = read_struct('8s',self.data,o)
  21. self.volume_space_lsb, o = read_struct('<I',self.data,o)
  22. self.volume_space_msb, o = read_struct('>I',self.data,o)
  23. _unused2, o = read_struct('32s',self.data,o)
  24. self.volume_set_lsb, o = read_struct('<H',self.data,o)
  25. self.volume_set_msb, o = read_struct('>H',self.data,o)
  26. self.volume_seq_lsb, o = read_struct('<H',self.data,o)
  27. self.volume_seq_msb, o = read_struct('>H',self.data,o)
  28. self.logical_block_size_lsb, o = read_struct('<H',self.data,o)
  29. self.logical_block_size_msb, o = read_struct('>H',self.data,o)
  30. self.path_table_size_lsb, o = read_struct('<I',self.data,o)
  31. self.path_table_size_msb, o = read_struct('>I',self.data,o)
  32. self.path_table_lsb, o = read_struct('<I',self.data,o)
  33. self.optional_path_table_lsb, o = read_struct('<I',self.data,o)
  34. self.path_table_msb, o = read_struct('>I',self.data,o)
  35. self.optional_path_table_msb, o = read_struct('>I',self.data,o)
  36. _offset = o
  37. self.root_dir_entry, o = read_struct('34s',self.data,o)
  38. self.root = ISOFile(self,_offset)
  39. self._cache = {}
  40. def get_file(self, path):
  41. if path == '/':
  42. return self.root
  43. else:
  44. if path in self._cache:
  45. return self._cache[path]
  46. units = path.split('/')
  47. units = units[1:] # remove root
  48. me = self.root
  49. for i in units:
  50. next_file = me.find(i)
  51. if not next_file:
  52. me = None
  53. break
  54. else:
  55. me = next_file
  56. self._cache[path] = me
  57. return me
  58. class ISOFile(object):
  59. def __init__(self, iso, offset):
  60. self.iso = iso
  61. self.offset = offset
  62. o = offset
  63. self.length, o = read_struct('B', self.iso.data, o)
  64. if not self.length:
  65. return
  66. self.ext_length, o = read_struct('B', self.iso.data, o)
  67. self.extent_start_lsb, o = read_struct('<I',self.iso.data, o)
  68. self.extent_start_msb, o = read_struct('>I',self.iso.data, o)
  69. self.extent_length_lsb, o = read_struct('<I',self.iso.data, o)
  70. self.extent_length_msb, o = read_struct('>I',self.iso.data, o)
  71. self.date_data, o = read_struct('7s', self.iso.data, o)
  72. self.flags, o = read_struct('b', self.iso.data, o)
  73. self.interleave_units, o = read_struct('b', self.iso.data, o)
  74. self.interleave_gap, o = read_struct('b', self.iso.data, o)
  75. self.volume_seq_lsb, o = read_struct('<H',self.iso.data, o)
  76. self.volume_seq_msb, o = read_struct('>H',self.iso.data, o)
  77. self.name_len, o = read_struct('b', self.iso.data, o)
  78. self.name, o = read_struct('{}s'.format(self.name_len), self.iso.data, o)
  79. self.name = self.name.decode('ascii')
  80. def write_extents(self):
  81. struct.pack_into('<I', self.iso.data, self.offset + 2, self.extent_start_lsb)
  82. struct.pack_into('>I', self.iso.data, self.offset + 6, self.extent_start_lsb)
  83. struct.pack_into('<I', self.iso.data, self.offset + 10, self.extent_length_lsb)
  84. struct.pack_into('>I', self.iso.data, self.offset + 14, self.extent_length_lsb)
  85. def readable_name(self):
  86. if not ';' in self.name:
  87. return self.name.lower()
  88. else:
  89. tmp, _ = self.name.split(';')
  90. return tmp.lower()
  91. def list(self):
  92. sectors = self.iso.data[self.extent_start_lsb * self.iso.sector_size: self.extent_start_lsb * self.iso.sector_size+ 3 * self.iso.sector_size]
  93. offset = 0
  94. while 1:
  95. f = ISOFile(self.iso, self.extent_start_lsb * self.iso.sector_size + offset)
  96. yield f
  97. offset += f.length
  98. if not f.length:
  99. break
  100. def find(self, name):
  101. sectors = self.iso.data[self.extent_start_lsb * self.iso.sector_size: self.extent_start_lsb * self.iso.sector_size+ 3 * self.iso.sector_size]
  102. offset = 0
  103. if '.' in name and len(name.split('.')[0]) > 8:
  104. a, b = name.split('.')
  105. name = a[:8] + '.' + b
  106. if '-' in name:
  107. name = name.replace('-','_')
  108. while 1:
  109. f = ISOFile(self.iso, self.extent_start_lsb * self.iso.sector_size + offset)
  110. if not f.length:
  111. if offset < self.extent_length_lsb:
  112. offset += 1
  113. continue
  114. else:
  115. break
  116. if ';' in f.name:
  117. tmp, _ = f.name.split(';')
  118. if tmp.endswith('.'):
  119. tmp = tmp[:-1]
  120. if tmp.lower() == name.lower():
  121. return f
  122. elif f.name.lower() == name.lower():
  123. return f
  124. offset += f.length
  125. return None
  126. class FAT(object):
  127. def __init__(self, iso, offset):
  128. self.iso = iso
  129. self.offset = offset
  130. self.bytespersector, _ = read_struct('H', self.iso.data, offset + 11)
  131. self.sectorspercluster, _ = read_struct('B', self.iso.data, offset + 13)
  132. self.reservedsectors, _ = read_struct('H', self.iso.data, offset + 14)
  133. self.numberoffats, _ = read_struct('B', self.iso.data, offset + 16)
  134. self.numberofdirs, _ = read_struct('H', self.iso.data, offset + 17)
  135. self.fatsize, _ = read_struct('H', self.iso.data, offset + 22)
  136. self.root_dir_sectors = (self.numberofdirs * 32 + (self.bytespersector - 1)) // self.bytespersector
  137. self.first_data_sector = self.reservedsectors + (self.numberoffats * self.fatsize) + self.root_dir_sectors
  138. self.root_sector= self.first_data_sector - self.root_dir_sectors
  139. self.root = FATDirectory(self, self.offset + self.root_sector * self.bytespersector)
  140. def get_offset(self, cluster):
  141. return self.offset + ((cluster - 2) * self.sectorspercluster + self.first_data_sector) * self.bytespersector
  142. def get_file(self, path):
  143. units = path.split('/')
  144. units = units[1:]
  145. me = self.root
  146. out = None
  147. for i in units:
  148. for fatfile in me.list():
  149. if fatfile.readable_name() == i:
  150. me = fatfile.to_dir()
  151. out = fatfile
  152. break
  153. else:
  154. return None
  155. return out
  156. class FATDirectory(object):
  157. def __init__(self, fat, offset):
  158. self.fat = fat
  159. self.offset = offset
  160. def list(self):
  161. o = self.offset
  162. while 1:
  163. out = FATFile(self.fat, o)
  164. if out.name != '\0\0\0\0\0\0\0\0':
  165. yield out
  166. else:
  167. break
  168. o += out.size
  169. class FATFile(object):
  170. def __init__(self, fat, offset):
  171. self.fat = fat
  172. self.offset = offset
  173. self.magic_long = None
  174. self.size = 0
  175. self.long_name = ''
  176. o = self.offset
  177. self.actual_offset = o
  178. self.attrib, _ = read_struct('B',self.fat.iso.data,o+11)
  179. while (self.attrib & 0x0F) == 0x0F:
  180. # Long file name entry
  181. tmp = read_struct('10s',self.fat.iso.data,o+1)[0]
  182. tmp += read_struct('12s',self.fat.iso.data,o+14)[0]
  183. tmp += read_struct('4s',self.fat.iso.data,o+28)[0]
  184. tmp = "".join([chr(x) for x in tmp[::2] if x != '\xFF']).strip('\x00')
  185. self.long_name = tmp + self.long_name
  186. self.size += 32
  187. o = self.offset + self.size
  188. self.actual_offset = o
  189. self.attrib, _ = read_struct('B',self.fat.iso.data,o+11)
  190. o = self.offset + self.size
  191. self.name, o = read_struct('8s',self.fat.iso.data,o)
  192. self.ext, o = read_struct('3s',self.fat.iso.data,o)
  193. self.attrib, o = read_struct('B',self.fat.iso.data,o)
  194. self.userattrib, o = read_struct('B',self.fat.iso.data,o)
  195. self.undelete, o = read_struct('b',self.fat.iso.data,o)
  196. self.createtime, o = read_struct('H',self.fat.iso.data,o)
  197. self.createdate, o = read_struct('H',self.fat.iso.data,o)
  198. self.accessdate, o = read_struct('H',self.fat.iso.data,o)
  199. self.clusterhi, o = read_struct('H',self.fat.iso.data,o)
  200. self.modifiedti, o = read_struct('H',self.fat.iso.data,o)
  201. self.modifiedda, o = read_struct('H',self.fat.iso.data,o)
  202. self.clusterlow, o = read_struct('H',self.fat.iso.data,o)
  203. self.filesize, o = read_struct('I',self.fat.iso.data,o)
  204. self.name = self.name.decode('ascii')
  205. self.ext = self.ext.decode('ascii')
  206. self.size += 32
  207. self.cluster = (self.clusterhi << 16) + self.clusterlow
  208. def is_dir(self):
  209. return bool(self.attrib & 0x10)
  210. def is_long(self):
  211. return bool((self.attrib & 0x0F) == 0x0F)
  212. def to_dir(self):
  213. return FATDirectory(self.fat, self.fat.get_offset(self.cluster))
  214. def get_offset(self):
  215. return self.fat.get_offset(self.cluster)
  216. def readable_name(self):
  217. if self.long_name:
  218. return self.long_name
  219. if self.ext.strip():
  220. return (self.name.strip() + '.' + self.ext.strip()).lower()
  221. else:
  222. return self.name.strip().lower()
  223. image = ISO('image.iso')
  224. fat = image.root.find('FAT.IMG')
  225. fatfs = FAT(image, fat.extent_start_lsb * image.sector_size)
  226. def process(fatfile, path):
  227. if fatfile.is_long():
  228. return
  229. if fatfile.readable_name() == '.':
  230. return
  231. if fatfile.readable_name() == '..':
  232. return
  233. if fatfile.is_dir():
  234. for i in fatfile.to_dir().list():
  235. process(i, path + fatfile.readable_name() + '/')
  236. else:
  237. cdfile = image.get_file(path + fatfile.readable_name())
  238. if not cdfile:
  239. if fatfile.readable_name() != 'bootia32.efi' and fatfile.readable_name() != 'bootx64.efi':
  240. print("Warning:", fatfile.readable_name(), "not found in ISO")
  241. else:
  242. cdfile.extent_start_lsb = fatfile.get_offset() // 2048
  243. cdfile.extent_length_lsb = fatfile.filesize
  244. cdfile.write_extents()
  245. for i in fatfs.root.list():
  246. process(i,'/')
  247. with open('image.iso','wb') as f:
  248. f.write(image.data)