summaryrefslogtreecommitdiffstats
path: root/libpwman/fileobj.py
blob: 3f09afb708ad0e874db58c60bc04bb1a110eacd0 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
# -*- coding: utf-8 -*-
"""
# Simple object file format.
# Copyright (c) 2011-2019 Michael Buesch <m@bues.ch>
# Licensed under the GNU/GPL version 2 or later.
"""

import errno

__all__ = [
	"FileObjError",
	"FileObj",
	"FileObjCollection",
]

class FileObjError(Exception):
	pass

class FileObj(object):
	# Raw object layout:
	#   [ 1 byte  ] => Name length
	#   [ x bytes ] => Name
	#   [ 4 bytes ] => Payload data length
	#   [ x bytes ] => Payload data

	def __init__(self, name, data):
		"""Construct FileObj().
		name: The object name. Must be bytes-like.
		data: The object payload. Must be bytes-like.
		"""
		assert isinstance(name, (bytes, bytearray)),\
		       "FileObj: Invalid 'name' type."
		assert isinstance(data, (bytes, bytearray)),\
		       "FileObj: Invalid 'data' type."
		if len(name) > 0xFF:
			raise FileObjError("FileObj: Name too long")
		self.__name = name
		if len(data) > 0xFFFFFFFF:
			raise FileObjError("FileObj: Data too long")
		self.__data = data

	def getName(self):
		return self.__name

	def getData(self):
		return self.__data

	def getRaw(self):
		r = bytearray()
		nameLen = len(self.__name)
		r += b"%c" % (nameLen & 0xFF)
		r += self.__name
		dataLen = len(self.__data)
		r += b"%c" % (dataLen & 0xFF)
		r += b"%c" % ((dataLen >> 8) & 0xFF)
		r += b"%c" % ((dataLen >> 16) & 0xFF)
		r += b"%c" % ((dataLen >> 24) & 0xFF)
		r += self.__data
		return r

	@classmethod
	def parseRaw(cls, raw):
		assert isinstance(raw, (bytes, bytearray)),\
		       "FileObj: Invalid 'raw' type."
		try:
			off = 0
			nameLen = raw[off]
			off += 1
			name = raw[off : off + nameLen]
			off += nameLen
			dataLen = (raw[off] |
				   (raw[off + 1] << 8) |
				   (raw[off + 2] << 16) |
				   (raw[off + 3] << 24))
			off += 4
			data = raw[off : off + dataLen]
			off += dataLen
		except (IndexError, KeyError) as e:
			raise FileObjError("Failed to parse file object")
		return (cls(name, data),
			off)

class FileObjCollection(object):
	def __init__(self, *objects):
		self.objects = objects

	def writeFile(self, filepath):
		try:
			with open(filepath, "wb") as f:
				f.write(self.getRaw())
				f.flush()
		except IOError as e:
			raise FileObjError("Failed to write file: %s" %
					   e.strerror)

	def getRaw(self):
		raw = bytearray()
		for obj in self.objects:
			raw += obj.getRaw()
		return raw

	def get(self, name):
		return [ o.getData()
			 for o in self.objects
			 if o.getName() == name ]

	def getOne(self, name, errorMsg=None, default=None):
		objs = self.get(name)
		if len(objs) != 1:
			if errorMsg:
				raise FileObjError(errorMsg)
			return default
		return objs[0]

	@classmethod
	def parseRaw(cls, raw):
		assert isinstance(raw, (bytes, bytearray)),\
		       "FileObjCollection: Invalid 'raw' type."
		offset = 0
		objects = []
		while offset < len(raw):
			(obj, objLen) = FileObj.parseRaw(raw[offset:])
			objects.append(obj)
			offset += objLen
		return cls(*objects)

	@classmethod
	def parseFile(cls, filepath):
		try:
			with open(filepath, "rb") as f:
				rawData = f.read()
		except (IOError) as e:
			if e.errno != errno.ENOENT:
				raise FileObjError("Failed to read file: %s" %\
						   e.strerror)
			return None
		return cls.parseRaw(rawData)
bues.ch cgit interface