From 62c1abfd24539dd25249b2a4ac27b62f239b716d Mon Sep 17 00:00:00 2001
From: acgnhiki <acgnhiki@outlook.com>
Date: Sat, 20 Nov 2021 10:38:49 +0800
Subject: [PATCH] release: v1.2.1

---
 CHANGELOG.md                      |  4 ++++
 README.md                         |  6 ++++++
 src/blrec/__init__.py             |  2 +-
 src/blrec/flv/data_analyser.py    |  6 ++++--
 src/blrec/flv/models.py           | 14 +++++++++++++-
 src/blrec/flv/stream_processor.py | 12 ++++++++++--
 6 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3adbdf3..2386cb1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # 更新日志
 
+## 1.2.1
+
+- 兼容 flv 头不正确的直播流
+
 ## 1.2.0
 
 - 改进文件处理方式，文件录制完成后就进行处理。
diff --git a/README.md b/README.md
index 9e3d6eb..f800635 100644
--- a/README.md
+++ b/README.md
@@ -146,8 +146,14 @@
 
 [CHANGELOG](CHANGELOG.md)
 
+---
+
 ## 赞助 & 支持
 
 如果觉得这个工具好用，对你有所帮助，可以投喂支持亿下哦~
 
 投喂赞助 ☞ <https://afdian.net/@acgnhiki>
+
+## Thanks
+
+[![JetBrains Logo (Main) logo](https://resources.jetbrains.com/storage/products/company/brand/logos/jb_beam.svg)](https://jb.gg/OpenSource)
diff --git a/src/blrec/__init__.py b/src/blrec/__init__.py
index 3acca5a..73cfba1 100644
--- a/src/blrec/__init__.py
+++ b/src/blrec/__init__.py
@@ -1,4 +1,4 @@
 
 __prog__ = 'blrec'
-__version__ = '1.2.0'
+__version__ = '1.2.1'
 __github__ = 'https://github.com/acgnhiki/blrec'
diff --git a/src/blrec/flv/data_analyser.py b/src/blrec/flv/data_analyser.py
index b05c919..1136106 100644
--- a/src/blrec/flv/data_analyser.py
+++ b/src/blrec/flv/data_analyser.py
@@ -75,6 +75,8 @@ class DataAnalyser:
         self._keyframe_filepositions: List[int] = []
         self._resolution: Optional[Resolution] = None
 
+        self._has_audio = False
+        self._has_video = False
         self._header_analysed = False
         self._audio_analysed = False
         self._video_analysed = False
@@ -83,8 +85,6 @@ class DataAnalyser:
         assert not self._header_analysed
         self._header_analysed = True
         self._size_of_flv_header = header.size
-        self._has_audio = header.has_audio()
-        self._has_video = header.has_video()
 
     def analyse_tag(self, tag: FlvTag) -> None:
         if is_audio_tag(tag):
@@ -198,6 +198,7 @@ class DataAnalyser:
 
     def _analyse_audio_tag(self, tag: AudioTag) -> None:
         if not self._audio_analysed:
+            self._has_audio = True
             self._audio_analysed = True
             self._audio_codec_id = tag.sound_format.value
             self._audio_sample_rate = tag.sound_rate.value
@@ -219,6 +220,7 @@ class DataAnalyser:
             pass
 
         if not self._video_analysed:
+            self._has_video = True
             self._video_analysed = True
             self._video_codec_id = tag.codec_id.value
 
diff --git a/src/blrec/flv/models.py b/src/blrec/flv/models.py
index 5941682..954fc8b 100644
--- a/src/blrec/flv/models.py
+++ b/src/blrec/flv/models.py
@@ -81,7 +81,7 @@ def non_negative_integer_validator(instance, attribute, value):  # type: ignore
         raise ValueError(f"'{attribute}' has to be a non negative integer!")
 
 
-@attr.s(auto_attribs=True, slots=True, frozen=True)
+@attr.s(auto_attribs=True, slots=True)
 class FlvHeader:
     signature: str
     version: int
@@ -94,6 +94,18 @@ class FlvHeader:
     def has_audio(self) -> bool:
         return bool(self.type_flag & 0b0000_0100)
 
+    def set_video_flag(self, value: bool) -> None:
+        if value:
+            self.type_flag |= 0b0000_0001
+        else:
+            self.type_flag &= ~0b0000_0001
+
+    def set_audio_flag(self, value: bool) -> None:
+        if value:
+            self.type_flag |= 0b0000_0100
+        else:
+            self.type_flag &= ~0b0000_0100
+
     @property
     def size(self) -> int:
         return self.data_offset
diff --git a/src/blrec/flv/stream_processor.py b/src/blrec/flv/stream_processor.py
index 2d3d042..5c5af8c 100644
--- a/src/blrec/flv/stream_processor.py
+++ b/src/blrec/flv/stream_processor.py
@@ -208,7 +208,6 @@ class StreamProcessor:
         self._in_reader = FlvReaderWithTimestampFix(stream)
         flv_header = self._read_header()
         self._has_audio = flv_header.has_audio()
-        assert flv_header.has_video(), 'no video in the stream!'
 
         try:
             first_data_tag = self._read_first_data_tag()
@@ -230,7 +229,7 @@ class StreamProcessor:
         self._new_file()
 
         try:
-            self._write_header(flv_header)
+            self._write_header(self._ensure_header_correct(flv_header))
             self._transfer_meta_tags()
             self._transfer_first_data_tag(first_data_tag)
         except Exception:
@@ -486,6 +485,15 @@ class StreamProcessor:
         self._size_updates.on_next(size)
         self._time_updates.on_next(tag.timestamp)
 
+    def _ensure_header_correct(self, header: FlvHeader) -> FlvHeader:
+        header.set_video_flag(
+            self._parameters_checker.last_video_header_tag is not None
+        )
+        header.set_audio_flag(
+            self._parameters_checker.last_audio_header_tag is not None
+        )
+        return header
+
     def _ensure_ts_correct(self, tag: FlvTag) -> None:
         if tag.timestamp + self._delta < 0:
             self._delta = -tag.timestamp