-
Notifications
You must be signed in to change notification settings - Fork 76
/
Copy pathrtp_h264_export.lua
344 lines (308 loc) · 13.8 KB
/
rtp_h264_export.lua
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
-- Dump RTP h.264 payload to raw h.264 file (*.264)
-- According to RFC3984 to dissector H264 payload of RTP to NALU, and write it
-- to from<sourceIp_sourcePort>to<dstIp_dstPort>.264 file. By now, we support single NALU,
-- STAP-A and FU-A format RTP payload for H.264.
-- You can access this feature by menu "Tools"
-- Author: Huang Qiangxiong ([email protected])
-- Modify by Yang Xing ([email protected])
------------------------------------------------------------------------------------------------
do
local version_str = string.match(_VERSION, "%d+[.]%d*")
local version_num = version_str and tonumber(version_str) or 5.1
-- lua>=5.4 直接使用位操作
-- 使用bit32进行位操作
if (version_num >= 5.4) then
local function band(a,b)
return(a&b)
end
local function bor(a,b)
return(a|b)
end
local function lshift(a,b)
return(a<<b)
end
else
local bit = (version_num >= 5.2) and require("bit32") or require("bit")
end
function string.starts(String,Start)
return string.sub(String,1,string.len(Start))==Start
end
function string.ends(String,End)
return End=='' or string.sub(String,-string.len(End))==End
end
function get_temp_path()
local tmp = nil
if tmp == nil or tmp == '' then
tmp = os.getenv('HOME')
if tmp == nil or tmp == '' then
tmp = os.getenv('USERPROFILE')
if tmp == nil or tmp == '' then
tmp = persconffile_path('temp')
else
tmp = tmp .. "/wireshark_temp"
end
else
tmp = tmp .. "/wireshark_temp"
end
end
return tmp
end
function get_ffmpeg_path()
local tmp = nil
if tmp == nil or tmp == '' then
tmp = os.getenv('FFMPEG')
if tmp == nil or tmp == '' then
tmp = ""
else
if not string.ends(tmp, "/bin/") then
tmp = tmp .. "/bin/"
end
end
end
return tmp
end
-- for geting h264 data (the field's value is type of ByteArray)
local f_h264 = Field.new("h264")
local f_rtp = Field.new("rtp")
local f_rtp_seq = Field.new("rtp.seq")
local f_rtp_timestamp = Field.new("rtp.timestamp")
local filter_string = nil
-- menu action. When you click "Tools->Export H264 to file" will run this function
local function export_h264_to_file()
-- window for showing information
local tw = TextWindow.new("Export H264 to File Info Win")
local pgtw;
-- add message to information window
function twappend(str)
tw:append(str)
tw:append("\n")
end
local ffmpeg_path = get_ffmpeg_path()
-- temp path
local temp_path = get_temp_path()
-- running first time for counting and finding sps+pps, second time for real saving
local first_run = true
local writed_nalu_begin = false
-- variable for storing rtp stream and dumping parameters
local stream_infos = nil
-- trigered by all h264 packats
local list_filter = ''
if filter_string == nil or filter_string == '' then
list_filter = "h264"
elseif string.find(filter_string,"h264")~=nil then
list_filter = filter_string
else
list_filter = "h264 && "..filter_string
end
twappend("Listener filter: " .. list_filter .. "\n")
local my_h264_tap = Listener.new("frame", list_filter)
-- get rtp stream info by src and dst address
function get_stream_info(pinfo)
local key = "from_" .. tostring(pinfo.src) .. "_" .. tostring(pinfo.src_port) .. "_to_" .. tostring(pinfo.dst) .. "_" .. tostring(pinfo.dst_port)
key = key:gsub(":", ".")
local stream_info = stream_infos[key]
if not stream_info then -- if not exists, create one
stream_info = { }
stream_info.filename = key.. ".264"
-- stream_info.filepath = stream_info.filename
-- stream_info.file,msg = io.open(stream_info.filename, "wb")
if not Dir.exists(temp_path) then
Dir.make(temp_path)
end
stream_info.filepath = temp_path.."/"..stream_info.filename
stream_info.file,msg = io.open(temp_path.."/"..stream_info.filename, "wb")
if msg then
twappend("io.open "..stream_info.filepath..", error "..msg)
end
-- twappend("Output file path:" .. stream_info.filepath)
stream_info.counter = 0 -- counting h264 total NALUs
stream_info.counter2 = 0 -- for second time running
stream_infos[key] = stream_info
twappend("Ready to export H.264 data (RTP from " .. tostring(pinfo.src) .. ":" .. tostring(pinfo.src_port)
.. " to " .. tostring(pinfo.dst) .. ":" .. tostring(pinfo.dst_port) .. " write to file:[" .. stream_info.filename .. "] ...")
end
return stream_info
end
-- write a NALU or part of NALU to file.
local function write_to_file(stream_info, str_bytes, begin_with_nalu_hdr, end_of_nalu)
if first_run then
stream_info.counter = stream_info.counter + 1
if begin_with_nalu_hdr then
-- save SPS PPS
local nalu_type = bit.band(str_bytes:byte(0,1), 0x1F)
if not stream_info.sps and nalu_type == 7 then
stream_info.sps = str_bytes
elseif not stream_info.pps and nalu_type == 8 then
stream_info.pps = str_bytes
end
end
else -- second time running
if not writed_nalu_begin then
if begin_with_nalu_hdr then
writed_nalu_begin = true
else
return
end
end
if stream_info.counter2 == 0 then
local nalu_type = bit.band(str_bytes:byte(0,1), 0x1F)
if nalu_type ~= 7 then
-- write SPS and PPS to file header first
if stream_info.sps then
stream_info.file:write("\x00\x00\x00\x01")
stream_info.file:write(stream_info.sps)
else
twappend("Not found SPS for [" .. stream_info.filename .. "], it might not be played!")
end
if stream_info.pps then
stream_info.file:write("\x00\x00\x00\x01")
stream_info.file:write(stream_info.pps)
else
twappend("Not found PPS for [" .. stream_info.filename .. "], it might not be played!")
end
end
end
if begin_with_nalu_hdr then
-- *.264 raw file format seams that every nalu start with 0x00000001
stream_info.file:write("\x00\x00\x00\x01")
end
stream_info.file:write(str_bytes)
stream_info.counter2 = stream_info.counter2 + 1
-- update progress window's progress bar
if stream_info.counter > 0 and stream_info.counter2 < stream_info.counter then
pgtw:update(stream_info.counter2 / stream_info.counter)
end
end
end
-- read RFC3984 about single nalu/stap-a/fu-a H264 payload format of rtp
-- single NALU: one rtp payload contains only NALU
local function process_single_nalu(stream_info, h264)
write_to_file(stream_info, h264:tvb():raw(), true, true)
end
-- STAP-A: one rtp payload contains more than one NALUs
local function process_stap_a(stream_info, h264)
local h264tvb = h264:tvb()
local offset = 1
repeat
local size = h264tvb(offset,2):uint()
write_to_file(stream_info, h264tvb:raw(offset+2, size), true, true)
offset = offset + 2 + size
until offset >= h264tvb:len()
end
-- FU-A: one rtp payload contains only one part of a NALU (might be begin, middle and end part of a NALU)
local function process_fu_a(stream_info, h264)
local h264tvb = h264:tvb()
local fu_idr = h264:get_index(0)
local fu_hdr = h264:get_index(1)
local end_of_nalu = (bit.band(fu_hdr, 0x40) ~= 0)
if bit.band(fu_hdr, 0x80) ~= 0 then
-- start bit is set then save nalu header and body
local nalu_hdr = bit.bor(bit.band(fu_idr, 0xE0), bit.band(fu_hdr, 0x1F))
write_to_file(stream_info, string.char(nalu_hdr) .. h264tvb:raw(2), true, end_of_nalu)
else
-- start bit not set, just write part of nalu body
write_to_file(stream_info, h264tvb:raw(2), false, end_of_nalu)
end
end
-- call this function if a packet contains h264 payload
function my_h264_tap.packet(pinfo,tvb)
if stream_infos == nil then
-- not triggered by button event, so do nothing.
return
end
local h264s = { f_h264() } -- using table because one packet may contains more than one RTP
for i,h264_f in ipairs(h264s) do
if h264_f.len < 2 then
return
end
local h264 = h264_f.range:bytes()
local hdr_type = bit.band(h264:get_index(0), 0x1F)
local stream_info = get_stream_info(pinfo)
if hdr_type > 0 and hdr_type < 24 then
-- Single NALU
process_single_nalu(stream_info, h264)
elseif hdr_type == 24 then
-- STAP-A Single-time aggregation
process_stap_a(stream_info, h264)
elseif hdr_type == 28 then
-- FU-A
process_fu_a(stream_info, h264)
else
twappend("Error: No.=" .. tostring(pinfo.number) .. " unknown type=" .. hdr_type .. " ; we only know 1-23(Single NALU),24(STAP-A),28(FU-A)!")
end
end
end
-- close all open files
local function close_all_files()
twappend("")
local index = 0;
if stream_infos then
local no_streams = true
for id,stream in pairs(stream_infos) do
if stream and stream.file then
stream.file:flush()
stream.file:close()
stream.file = nil
index = index + 1
twappend(index .. ": [" .. stream.filename .. "] generated OK!")
local anony_fuc = function ()
twappend("ffplay -x 640 -y 640 -autoexit "..stream.filename)
--copy_to_clipboard("ffplay -x 640 -y 640 -autoexit "..stream.filepath)
os.execute(ffmpeg_path.."ffplay -x 640 -y 640 -autoexit "..stream.filepath)
end
tw:add_button("Play "..index, anony_fuc)
no_streams = false
end
end
if no_streams then
twappend("Not found any H.264 over RTP streams!")
else
tw:add_button("Browser", function () browser_open_data_file(temp_path) end)
end
end
end
function my_h264_tap.reset()
-- do nothing now
end
tw:set_atclose(function ()
my_h264_tap:remove()
if Dir.exists(temp_path) then
Dir.remove_all(temp_path)
end
end)
local function export_h264()
pgtw = ProgDlg.new("Export H264 to File Process", "Dumping H264 data to file...")
first_run = true
stream_infos = {}
-- first time it runs for counting h.264 packets and finding SPS and PPS
retap_packets()
first_run = false
-- second time it runs for saving h264 data to target file.
retap_packets()
close_all_files()
-- close progress window
pgtw:close()
stream_infos = nil
end
tw:add_button("Export All", function ()
export_h264()
end)
tw:add_button("Set Filter", function ()
tw:close()
dialog_menu()
end)
end
local function dialog_func(str)
filter_string = str
export_h264_to_file()
end
function dialog_menu()
new_dialog("Filter Dialog",dialog_func,"Filter")
end
local function dialog_default()
filter_string = get_filter()
export_h264_to_file()
end
-- Find this feature in menu "Tools"
register_menu("Video/Export H264", dialog_default, MENU_TOOLS_UNSORTED)
end